b87368e1e7fd832b505db9cc08015ac7af8f95de,VAE/main.py,,train,#Any#,94

Before Change

def train(epoch):
    batch_data_t = torch.FloatTensor(BATCH_SIZE, 784)
    if cuda:
        batch_data_t = batch_data_t.cuda()
    batch_data = Variable(batch_data_t, requires_grad=False)
    for i in range(0, training_data.size(0), BATCH_SIZE):

After Change

    model.train()
    train_loss = 0
    for batch in train_loader:
        batch = Variable(batch)

        optimizer.zero_grad()
        recon_batch, mu, logvar = model(batch)
        loss = loss_function(recon_batch, batch, mu, logvar)

Instances