fix beta being a tensor in post-hoc ema

lucidrains · Aug 27, 2024 · 1aa934c · 1aa934c
1 parent d637fdc
commit 1aa934c
Show file tree

Hide file tree

Showing 2 changed files with 2 additions and 2 deletions.
diff --git a/ema_pytorch/post_hoc_ema.py b/ema_pytorch/post_hoc_ema.py
@@ -123,7 +123,7 @@ def model(self):
 
     @property
     def beta(self):
-        return (1 - 1 / (self.step + 1)) ** (1 + self.gamma)
+        return (1. - 1. / (self.step.item() + 1.)) ** (1. + self.gamma)
 
     def eval(self):
         return self.ema_model.eval()

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