🔧 [Update] Trainer input config

examples/example_train.py

@@ -28,7 +28,7 @@ def main(cfg: Config):
     # TODO: get_device or rank, for DDP mode
     device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 
-    trainer = Trainer(model, cfg.hyper.train, device)
+    trainer = Trainer(model, cfg, device)
     trainer.train(dataloader, 10)
 
 

yolo/tools/trainer.py

@@ -9,7 +9,9 @@ from yolo.utils.loss import get_loss_function
 
 
 class Trainer:
-    def __init__(self, model: YOLO, train_cfg: TrainConfig, device):
+    def __init__(self, model: YOLO, cfg: Config, device):
+        train_cfg: TrainConfig = cfg.hyper.train
+
         self.model = model.to(device)
         self.device = device
         self.optimizer = get_optimizer(model.parameters(), train_cfg.optimizer)