Fix ema resume checkpoint (#184)

neuralmagic · Feb 28, 2023 · bd2bc0b · bd2bc0b
1 parent 25f226b
commit bd2bc0b
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/utils/neuralmagic/sparsification_manager.py b/utils/neuralmagic/sparsification_manager.py
@@ -601,7 +601,7 @@ def maybe_create_sparsification_manager(
 
         # reconstruct ToggleableModelEMA from state dictionary
         if ckpt["ema"]:
-            ckpt["ema"] = load_ema(ckpt["ema"], model)
+            ckpt["ema"] = load_ema(ckpt["ema"], model).ema
 
         return sparsification_manager