From 735956ee5afb1d51e96106efecca422b5b411ee5 Mon Sep 17 00:00:00 2001 From: Rahul Tuli Date: Wed, 11 Jan 2023 09:48:32 -0500 Subject: [PATCH] bugfix, set loss --- src/sparseml/pytorch/torchvision/train.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/sparseml/pytorch/torchvision/train.py b/src/sparseml/pytorch/torchvision/train.py index 04b7eb3190a..56f6ea0f141 100644 --- 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,