Fix emulated step condition

neuralmagic · Jan 30, 2023 · 53bcf5b · 53bcf5b
1 parent 59efee0
commit 53bcf5b
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/train.py b/train.py
@@ -386,7 +386,7 @@ def _create_dataloaders():
                     ema.update(model)
                 last_opt_step = ni
 
-            elif sparsification_manager:
+            elif sparsification_manager and sparsification_manager.train_manager:
                 # Call for SparseML integration since the number of steps per epoch can vary
                 # This keeps the number of steps per epoch equivalent to the number of batches per epoch
                 # Does not step the scaler or the optimizer