Torchvision scaler (#1312)

* Passing scaler into torchvision manager.modify

* Properly wrapping optim/scaler for amp

* Properly disabling scaler

Author: corey-nm

src/sparseml/pytorch/torchvision/train.py

@@ -574,9 +574,17 @@ def log_metrics(tag: str, metrics: utils.MetricLogger, epoch: int, epoch_step: i
             loggers=logger,
             distillation_teacher=args.distill_teacher,
         )
-        optimizer = manager.modify(
-            model, optimizer, steps_per_epoch=steps_per_epoch, epoch=args.start_epoch
+        step_wrapper = manager.modify(
+            model,
+            optimizer,
+            steps_per_epoch=steps_per_epoch,
+            epoch=args.start_epoch,
+            wrap_optim=scaler,
         )
+        if scaler is None:
+            optimizer = step_wrapper
+        else:
+            scaler = step_wrapper
 
     lr_scheduler = _get_lr_scheduler(
         args, optimizer, checkpoint=checkpoint, manager=manager
@@ -597,7 +605,7 @@ def log_metrics(tag: str, metrics: utils.MetricLogger, epoch: int, epoch_step: i
         if args.distributed:
             train_sampler.set_epoch(epoch)
         if manager is not None and manager.qat_active(epoch=epoch):
-            scaler = None
+            scaler._enabled = False
             model_ema = None
 
         train_metrics = train_one_epoch(