def train_one_epoch(train_dataloader, model, optimizer, criterion, opt, epoch):
    """
    This function is called every epoch to perform training of the network over one full
    (randomized) iteration of the dataset.

    Args:
        train_dataloader: torch.utils.data.DataLoader, returns (augmented) training data.
        model:            Network to train.
        optimizer:        Optimizer to use for training.
        criterion:        criterion to use during training.
        opt:              argparse.Namespace, Contains all relevant parameters.
        epoch:            int, Current epoch.

    Returns:
        Nothing!
    """

    loss_collect = []

    start = time.time()
    data_iterator = tqdm(
        train_dataloader, desc="Epoch {} Training...".format(epoch)
    )

    for i, (class_labels, input) in enumerate(data_iterator):
        # Compute embeddings for input batch
        features = model(input.to(opt.device))

        # Compute loss.
        if opt.loss != "smoothap":
            loss = criterion(features, class_labels)
        else:
            loss = criterion(features)

        # Ensure gradients are set to zero at beginning
        optimizer.zero_grad()
        # Compute gradient
        loss.backward()

        train_dataloader.dataset.classes_visited = []

        if opt.grad_measure:
            # If desired, save computed gradients.
            grad_measure.include(model.model.last_linear)

        # Update weights using comp. gradients.
        optimizer.step()

        # Store loss per iteration.
        loss_collect.append(loss.item())
        if i == len(train_dataloader) - 1:
            data_iterator.set_description(
                "Epoch (Train) {0}: Mean Loss [{1:.4f}]".format(
                    epoch, np.mean(loss_collect)
                )
            )

    # Save metrics
    LOG.log(
        "train",
        LOG.metrics_to_log["train"],
        [epoch, np.round(time.time() - start, 4), np.mean(loss_collect)],
    )
    writer.add_scalar("global/training_loss", np.mean(loss_collect), epoch)
    if opt.grad_measure:
        # Dump stored gradients to Pickle-File.
        grad_measure.dump(epoch)

    This function is called every epoch to perform training of the network over one full
    (randomized) iteration of the dataset.

    Args:
        train_dataloader: torch.utils.data.DataLoader, returns (augmented) training data.
        model:            Network to train.
        optimizer:        Optimizer to use for training.
        criterion:        criterion to use during training.
        opt:              argparse.Namespace, Contains all relevant parameters.
        epoch:            int, Current epoch.

    Returns:
        Nothing!
    
examples/pytorch/hilander/PSS/Smooth_AP/src/main.py