address #28

lucidrains · Aug 7, 2024 · a6ac238 · a6ac238
1 parent c427d1b
commit a6ac238
Show file tree

Hide file tree

Showing 2 changed files with 6 additions and 2 deletions.
diff --git a/ema_pytorch/post_hoc_ema.py b/ema_pytorch/post_hoc_ema.py
@@ -304,10 +304,14 @@ def step(self):
     def device(self):
         return self.step.device
 
-    def copy_params_from_ema_to_model(self):
+    def copy_params_from_model_to_ema(self):
         for ema_model in self.ema_models:
             ema_model.copy_params_from_model_to_ema()
 
+    def copy_params_from_ema_to_model(self):
+        for ema_model in self.ema_models:
+            ema_model.copy_params_from_ema_to_model()
+
     def update(self):
         for ema_model in self.ema_models:
             ema_model.update()

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