store the ema model in f16 to save on space for post hoc ema

lucidrains · Feb 8, 2024 · c77058c · c77058c
1 parent 8ad8913
commit c77058c
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
@@ -320,7 +320,7 @@ def checkpoint(self):
             filename = f'{ind}.{step}.pt'
             path = self.checkpoint_folder / filename
 
-            pkg = ema_model.state_dict()
+            pkg = deepcopy(ema_model).half().state_dict()
             torch.save(pkg, str(path))
 
     @beartype

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