One last place where checkpoint_num_classes should be present (#1549)

Deci-AI · Oct 18, 2023 · 68a96bb · 68a96bb
1 parent 096cb0e
commit 68a96bb
Showing 1 changed file with 1 addition and 0 deletions.
diff --git a/src/super_gradients/training/sg_trainer/sg_trainer.py b/src/super_gradients/training/sg_trainer/sg_trainer.py
@@ -265,6 +265,7 @@ def train_from_config(cls, cfg: Union[DictConfig, dict]) -> Tuple[nn.Module, Tup
             pretrained_weights=cfg.checkpoint_params.pretrained_weights,
             checkpoint_path=cfg.checkpoint_params.checkpoint_path,
             load_backbone=cfg.checkpoint_params.load_backbone,
+            checkpoint_num_classes=get_param(cfg.checkpoint_params, "checkpoint_num_classes"),
         )
 
         # INSTANTIATE DATA LOADERS