Merge branch 'main' into torchvision-dp-support

neuralmagic · Jan 26, 2023 · f3550bb · f3550bb
2 parents 86f9575 + 8353866
commit f3550bb
Show file tree

Hide file tree

Showing 2 changed files with 6 additions and 5 deletions.
diff --git a/src/sparseml/pytorch/image_classification/export.py b/src/sparseml/pytorch/image_classification/export.py
@@ -387,9 +387,8 @@ def main(
 
     if recipe is not None:
         ScheduledModifierManager.from_yaml(recipe).apply_structure(model)
-
-    if checkpoint_path:
-        load_model(checkpoint_path, model, strict=True)
+        if checkpoint_path:
+            load_model(checkpoint_path, model, strict=True)
 
     if one_shot is not None:
         ScheduledModifierManager.from_yaml(file_path=one_shot).apply(module=model)

diff --git a/src/sparseml/pytorch/torchvision/train.py b/src/sparseml/pytorch/torchvision/train.py
@@ -384,7 +384,7 @@ def collate_fn(batch):
 
     if args.distill_teacher not in ["self", "disable", None]:
         _LOGGER.info("Instantiating teacher")
-        args.distill_teacher = _create_model(
+        distill_teacher = _create_model(
             arch_key=args.teacher_arch_key,
             local_rank=local_rank,
             pretrained=True,  # teacher is always pretrained
@@ -393,6 +393,8 @@ def collate_fn(batch):
             device=device,
             num_classes=num_classes,
         )
+    else:
+        distill_teacher = args.distill_teacher
 
     if args.distributed and args.sync_bn:
         model = torch.nn.SyncBatchNorm.convert_sync_batchnorm(model)
@@ -576,7 +578,7 @@ def log_metrics(tag: str, metrics: utils.MetricLogger, epoch: int, epoch_step: i
             model,
             epoch=args.start_epoch,
             loggers=logger,
-            distillation_teacher=args.distill_teacher,
+            distillation_teacher=distill_teacher,
         )
         step_wrapper = manager.modify(
             model,