init

trainer/trainer.py

 import torch
+import importlib
 
 from utils.utils import *
 
 
-def train_one_epoch(model, optimizer, data_loader, device, epoch, loss, loss_weights, warmup=False):
+def train_one_epoch(model, optimizer, data_loader, device, loss, loss_weights, warmup=False):
+    """
+
+    :param model: the PyTorch model that will be trained.
+    :param optimizer:
+    :param data_loader:
+    :param device:
+    :param loss:
+    :param loss_weights:
+    :param warmup:
+    :return:
+    """
     lr_scheduler = None
     log_bar = None
     if is_main_process():