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):
        optimizer.zero_grad()
        batch_data.data[:] = training_data[i:i + BATCH_SIZE]
        recon_batch_data, mu, logvar = model(batch_data)
        loss = loss_function(recon_batch_data, batch_data, mu, logvar)
        loss.backward()
        loss = loss.data[0]
        optimizer.step()
        if i % 10 == 0:
            print("Epoch: {} [{}/{} ({:.0f}%)]\tLoss: {:.4f}".format(
                epoch,
                i + BATCH_SIZE, training_data.size(0),
                float(i + BATCH_SIZE) / training_data.size(0) * 100,
                loss / BATCH_SIZE))


def test(epoch):
    test_loss = 0
    batch_data_t = torch.FloatTensor(TEST_BATCH_SIZE, 784)
    if cuda:

After Change

def train(epoch):
    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)
        loss.backward()
        train_loss += loss
        optimizer.step()

    print("====> Epoch: {} Loss: {:.4f}".format(
          epoch,
          train_loss.data[0] / training_data.size(0)))

Instances