address #33 (comment)

lucidrains · Nov 21, 2024 · 26b6de3 · 26b6de3
1 parent 7f6cd75
commit 26b6de3
Show file tree

Hide file tree

Showing 3 changed files with 3 additions and 3 deletions.
diff --git a/ema_pytorch/ema_pytorch.py b/ema_pytorch/ema_pytorch.py
@@ -101,7 +101,7 @@ def __init__(
 
         # handle callable returning ema module
 
-        if callable(ema_model):
+        if not isinstance(ema_model, Module) and callable(ema_model):
             ema_model = ema_model()
 
         # ema model

diff --git a/ema_pytorch/post_hoc_ema.py b/ema_pytorch/post_hoc_ema.py
@@ -77,7 +77,7 @@ def __init__(
 
         # handle callable returning ema module
 
-        if callable(ema_model):
+        if not isinstance(ema_model, Module) and callable(ema_model):
             ema_model = ema_model()
 
         # ema model

diff --git a/setup.py b/setup.py
@@ -3,7 +3,7 @@
 setup(
   name = 'ema-pytorch',
   packages = find_packages(exclude=[]),
-  version = '0.7.4',
+  version = '0.7.5',
   license='MIT',
   description = 'Easy way to keep track of exponential moving average version of your pytorch module',
   author = 'Phil Wang',