bugfix, set loss

neuralmagic · Jan 11, 2023 · 735956e · 735956e
1 parent 2b2af45
commit 735956e
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/src/sparseml/pytorch/torchvision/train.py b/src/sparseml/pytorch/torchvision/train.py
@@ -91,7 +91,7 @@ def train_one_epoch(
 
         if steps_accumulated % args.gradient_accum_steps == 0:
             if manager is not None:
-                manager.loss_update(
+                loss = manager.loss_update(
                     loss=loss,
                     module=model,
                     optimizer=optimizer,