# Forward pass on real MNIST
out_dis, hid = Net_D(data)
c1 = LogSoftmax()(Q_cat(hid))
loss_dis = mse(out_dis, ones) - torch.sum(targets * c1) / (torch.sum(targets) + 1e-3) # Loss for real MNIST
# Forward pass on generated MNIST
out_gen = Net_G(z)
out_dis, hid = Net_D(out_gen)
# Loss for generated MNIST
loss_dis = loss_dis + mse(out_dis, zeros)
loss_dis = loss_dis
# Zero gradient buffers for gen and Q_cat and backward pass
Net_D.zero_grad()
Q_cat.zero_grad()
loss_dis.backward(retain_graph = True) # We need PyTorch to retain the graph buffers so we can run backward again later
d_optim.step() # Apply the discriminator's update now since we have to delete its gradients later
# And backward pass and loss for generator and update
Net_G.zero_grad()
loss_gen = mse(out_dis, ones)
loss_gen.backward(retain_graph = True)
Net_D.zero_grad() # Don't want the gradients of the generator's objective in the discriminator
# Forward pass and loss for latent codes
loss_q = 0
c1 = LogSoftmax()(Q_cat(hid))
loss_q += nll(c1, torch.max(z_dict['cat'], dim = 1)[1])
