small helper method

lucidrains · Oct 20, 2024 · dee87fb · dee87fb
1 parent 55959a3
commit dee87fb
Show file tree

Hide file tree

Showing 2 changed files with 9 additions and 1 deletion.
diff --git a/ema_pytorch/ema_pytorch.py b/ema_pytorch/ema_pytorch.py
@@ -200,6 +200,14 @@ def model(self):
     def eval(self):
         return self.ema_model.eval()
 
+    @torch.no_grad()
+    def forward_eval(self, *args, **kwargs):
+        # handy function for invoking ema model with no grad + eval
+        training = self.ema_model.training
+        out = self.ema_model(*args, **kwargs)
+        self.ema_model.train(training)
+        return out
+
     def restore_ema_model_device(self):
         device = self.initted.device
         self.ema_model.to(device)

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