960518e2734a42f7da19dc3d4d49501004c037e1,mnist/main.py,,,#,54

Before Change

        x = self.relu(self.fc2(x))
        return self.softmax(x)

model = Net().cuda()

criterion = nn.ClassNLLCriterion()

// Training settings
BATCH_SIZE = 150
TEST_BATCH_SIZE = 1000
NUM_EPOCHS = 2

optimizer = optim.SGD(model, lr=1e-2, momentum=0.9)

def train(epoch):
    batch_data = Variable(torch.cuda.FloatTensor(BATCH_SIZE, 1, 28, 28), requires_grad=False)
    batch_targets = Variable(torch.cuda.FloatTensor(BATCH_SIZE), requires_grad=False)

After Change

training_data = training_set[0].view(-1, 1, 28, 28).div(255)
training_labels = training_set[1]
test_data = test_set[0].view(-1, 1, 28, ).div(255)
test_labels = test_set[1]

del training_set

Instances