From bd2bc0b193dca24ed6da403c68f3120b17cd039e Mon Sep 17 00:00:00 2001
From: Konstantin Gulin <66528950+KSGulin@users.noreply.github.com>
Date: Tue, 28 Feb 2023 09:33:07 -0800
Subject: [PATCH] Fix ema resume checkpoint (#184)

---
 utils/neuralmagic/sparsification_manager.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/utils/neuralmagic/sparsification_manager.py b/utils/neuralmagic/sparsification_manager.py
index 546e94476ec6..0e9fca2b769e 100644
--- 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