Variational_Autoencoder_Training.py

vae = VAE()
vae.to(device)
criterion = nn.MSELoss()
optimizer = optim.Adamax(vae.parameters(),lr = 1e-4)

l = None
for epoch in range(100):
  for i, data in enumerate(loader,0):
    inputs,classes = data
    inputs,classes = Variable(inputs.resize_(batch_size,784)).to(device),Variable(classes).to(device)
    optimizer.zero_grad()
    out = vae(inputs)
    ll = latent_loss(vae.z_mean,vae.z_sigma)
    loss = criterion(out,inputs)+ll
    loss.backward()
    optimizer.step()
    l = loss.item()
  print(epoch,l)