code: def train_one_epoch_finetune(
      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 gt labels...".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)
  
file_path: dgl/examples/pytorch/hilander/PSS/Smooth_AP/src/finetune_1head.py
imports: import warnings
  import argparse
  import datetime
  import os
  import random
  import matplotlib
  import numpy
  from pathlib import Path
  import auxiliaries
  import datasets
  import evaluate
  import losses
  import netlib
  import torch.multiprocessing
  from tensorboardX import SummaryWriter
  from tqdm import tqdm
  import time
  import time
label_desc: Use image node_id 2 for calling a global function with example usage: train_one_epoch_finetune(train_dataloader, model, optimizer, criterion, opt, epoch) without return types