update ema-pytorch

Author: lucidrains

alphafold3_pytorch/trainer.py

@@ -222,7 +222,9 @@ def __init__(
         checkpoint_folder: str = './checkpoints',
         overwrite_checkpoints: bool = False,
         fabric_kwargs: dict = dict(),
-        ema_kwargs: dict = dict()
+        ema_kwargs: dict = dict(
+            use_foreach = True
+        )
     ):
         super().__init__()
 

pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "alphafold3-pytorch"
-version = "0.1.47"
+version = "0.1.48"
 description = "Alphafold 3 - Pytorch"
 authors = [
     { name = "Phil Wang", email = "[email protected]" }
@@ -27,7 +27,7 @@ dependencies = [
     "biopython>=1.83",
     "einops>=0.8.0",
     "einx>=0.2.2",
-    "ema-pytorch>=0.4.8",
+    "ema-pytorch>=0.5.0",
     "environs",
     "frame-averaging-pytorch>=0.0.18",
     "jaxtyping>=0.2.28",

tests/test_trainer.py

@@ -149,7 +149,12 @@ def test_trainer():
         valid_every = 1,
         grad_accum_every = 2,
         checkpoint_every = 1,
-        overwrite_checkpoints = True
+        overwrite_checkpoints = True,
+        ema_kwargs = dict(
+            use_foreach = True,
+            update_after_step = 0,
+            update_every = 1
+        )
     )
 
     trainer()