Fixing return argument of _create_model for distillation (#1345)

Co-authored-by: Tuan Nguyen <tuan@neuralmagic.com>

Author: 2 people

src/sparseml/pytorch/torchvision/train.py

@@ -393,7 +393,7 @@ def collate_fn(batch):
 
     if args.distill_teacher not in ["self", "disable", None]:
         _LOGGER.info("Instantiating teacher")
-        distill_teacher = _create_model(
+        distill_teacher, _ = _create_model(
             arch_key=args.teacher_arch_key,
             local_rank=local_rank,
             pretrained=True,  # teacher is always pretrained