My entire training code would be like so
def train(train_loader, model, criterion, optimizer, epoch):
    train_loss = AverageMeter()
    model.train()

    for i, (inputs, labels) in enumerate(train_loader):
     
        
        N = inputs.size(0)
        inputs = torch.FloatTensor(inputs).to(device)
        labels = torch.LongTensor(labels).to(device)

        optimizer.zero_grad()
        outputs = model(inputs)

        loss = criterion(outputs, labels.to(device)) / N
        loss.backward()
        optimizer.step()

        train_loss.update(loss.data[0], N)