import torch


def train_epoch(model, train_loader, criterion, optimizer, device):
    """Train the model for 1 epoch
    Args:
        model: nn.Module
        train_loader: train DataLoader
        criterion: callable loss function
        optimizer: pytorch optimizer
        device: torch.device
    Returns
    -------
    Tuple[Float, Float]
        average train loss and average train accuracy for current epoch
    """

    train_losses = []
    train_corrects = []
    model.train()

    # Iterate over data.
    for inputs, labels in train_loader:
        inputs = inputs.to(device)
        labels = labels.to(device)

        # prediction
        outputs = model(inputs)

        # calculate loss
        _, preds = torch.max(outputs, 1)
        loss = criterion(outputs, labels)

        # backprop
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()

        # statistics
        train_losses.append(loss.item())
        train_corrects.append(torch.sum(preds == labels.data).item())

    return sum(train_losses)/len(train_losses), sum(train_corrects)/len(train_loader.dataset)