import torch


def train(train_data, model, optimizer, criterion, device):
    """
    Trains/updates the model for one epoch on the training dataset.

    Parameters:
        train_data (torch tensor): trainset
        model (torch.nn.module): Model to be trained
        optimizer (torch.optim.optimizer): optimizer instance like SGD or Adam
        criterion (torch.nn.modules.loss): loss function like CrossEntropyLoss
        device (string): cuda or cpu
    """

    # switch to train mode
    model.train()

    # iterate through the dataset loader
    i = 0
    losses = []
    for (inp, target) in train_data:
        # transfer inputs and targets to the GPU (if it is available)
        inp = inp.to(device)
        target = target.to(device)

        # compute output, i.e. the model forward
        output = model(inp)

        # calculate the loss
        loss = criterion(output, target)
        # print("loss", loss)
        losses.append(loss)

        print("loss {:.2}".format(loss))
        # compute gradient and do the SGD step
        # we reset the optimizer with zero_grad to "flush" former gradients
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()

    avg_loss = torch.mean(torch.stack(losses)).item()
    print("avg.loss {:.2}".format(avg_loss))
    return avg_loss