Dropping rolling loss

import torch
import torchvision
import torchvision.transforms as transforms
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim
import torchvision.models as models


# Parameters
gpu = torch.device('cuda')
nb_epochs = 30
nb_runs_per_model = 5
bs = 32
models_list = ['vgg16_bn', 'vgg16']


def prep_data(bs):
    train_transform = transforms.Compose([transforms.RandomHorizontalFlip(),
                                          transforms.ToTensor(),
                                          transforms.Normalize((.5, .5, .5), (.5, .5, .5))])

    test_transform = transforms.Compose([transforms.ToTensor(),
                                         transforms.Normalize((.5, .5, .5), (.5, .5, .5))])

    trainset = torchvision.datasets.CIFAR10(root = "../data_CIFAR10", train = True, download = False, transform = train_transform)
    trainloader = torch.utils.data.DataLoader(trainset, batch_size = bs, shuffle = True, drop_last = True)

    validset = torchvision.datasets.CIFAR10(root = "../data_CIFAR10", train = False, download = False, transform = test_transform)
    validloader = torch.utils.data.DataLoader(validset, batch_size = bs, shuffle = False)

    classes = ('plane', 'car', 'bird', 'cat', 'deer', 'dog', 'frog', 'horse', 'ship', 'truck')

    return trainloader, validloader, classes



def train_model(model, trainLoader, criterion, optimizer, epochs, gpu):
    train_losses = []

    for epoch in range(epochs):
        for i, data in enumerate(trainLoader, 0):
            inputs, labels = data[0].to(gpu), data[1].to(gpu)

            optimizer.zero_grad()
            outputs = model(inputs)
            loss = criterion(outputs, labels)
            loss.backward()
            optimizer.step()

            train_losses.append(loss.item())

            if(i%100==99):
                print("[%d, %5d] loss : %.5f" % (epoch + 1, i + 1, loss.item()))

    print("Finished training\n\n")
    return train_losses



def main():
    for name in models_list:
        for i in range(nb_runs_per_model):

            trainLoader, validLoader, classes = prep_data(bs)

            print("Training {}, run number {} / {} : \n".format(name, i+1, nb_runs_per_model))

            if name == 'vgg16':
                model = models.vgg16(num_classes = 10)
            elif name == 'vgg16_bn':
                model = models.vgg16_bn(num_classes = 10)

            model.to(gpu)
            criterion = nn.CrossEntropyLoss()
            optimizer = optim.SGD(model.parameters(), lr = 0.003)

            train_losses = train_model(model, trainLoader, criterion, optimizer, nb_epochs, gpu)

            with open("Results.py", 'a') as f:
                f.write("### {} run {}/{} : \n{}_{} = {} \n".format(name, i+1, nb_runs_per_model,
                                                                    name, i+1, train_losses))


if __name__ == '__main__':
    main()