fix an issue @amorehead is running into during training

lucidrains · lucidrains · commit 5e0002c21bb4 · 2024-09-27T19:23:46.000-07:00
diff --git a/alphafold3_pytorch/trainer.py b/alphafold3_pytorch/trainer.py
@@ -219,6 +219,7 @@ def __init__(
                 beta = ema_decay,
                 include_online_model = False,
                 allow_different_devices = True,
+                coerce_dtype = True,
                 **ema_kwargs
             )
 
diff --git a/pyproject.toml b/pyproject.toml
@@ -1,6 +1,6 @@
 [project]
 name = "alphafold3-pytorch"
-version = "0.5.45"
+version = "0.5.46"
 description = "Alphafold 3 - Pytorch"
 authors = [
     { name = "Phil Wang", email = "lucidrains@gmail.com" },
@@ -33,7 +33,7 @@ dependencies = [
     "CoLT5-attention>=0.11.0",
     "einops>=0.8.0",
     "einx>=0.2.2",
-    "ema-pytorch>=0.5.0",
+    "ema-pytorch>=0.6.3",
     "environs",
     "lion-pytorch>=0.2.2",
     "joblib",

Original file line number	Diff line number	Diff line change
`@@ -219,6 +219,7 @@ def __init__(`
`219`	`219`	`beta = ema_decay,`
`220`	`220`	`include_online_model = False,`
`221`	`221`	`allow_different_devices = True,`
	`222`	`+ coerce_dtype = True,`
`222`	`223`	`**ema_kwargs`
`223`	`224`	`)`
`224`	`225`