Training

train.py

#Train the model
iter = 0
for epoch in range(num_epochs):
    for i, (images, labels) in enumerate(train_loader):
        images = Variable(images)
        labels = Variable(labels)
        
        #Clear the gradients
        optimizer.zero_grad()
        
        #Forward propagation 
        outputs = model(images)      
        
        #Calculating loss with softmax to obtain cross entropy loss
        loss = criterion(outputs, labels)
        
        #Backward propation
        loss.backward()
        
        #Updating gradients
        optimizer.step()
        
        iter += 1
        
        #Total number of labels
        total = labels.size(0)
        
        #Obtaining predictions from max value
        _, predicted = torch.max(outputs.data, 1)
        
        #Calculate the number of correct answers
        correct = (predicted == labels).sum().item()
        
        #Print loss and accuracy
        if (i + 1) % 100 == 0:
            print('Epoch [{}/{}], Step [{}/{}], Loss: {:.4f}, Accuracy: {:.2f}%'
                  .format(epoch + 1, num_epochs, i + 1, len(train_loader), loss.item(),
                          (correct / total) * 100))