def mdn_loss(mu, sigma, pi, target):
log_probs = Normal(mu, sigma).log_prob(target.unsqueeze(-1))
log_probs = torch.logsumexp(log_probs + pi, -1)
return -log_probs.mean()
def forward(self, x, mu, std, pi):
log_probs = Normal(mu, std + 1e-9).log_prob(x.unsqueeze(-1))
log_probs = torch.logsumexp(log_probs + pi, -1)
return -log_probs.mean()
