allow for using AdamAtan2 and get rid of the epsilon

Author: lucidrains

alphafold3_pytorch/trainer.py

@@ -39,6 +39,8 @@
 from torch.utils.data import Sampler, Dataset, DataLoader as OrigDataLoader
 from torch.optim.lr_scheduler import LambdaLR, LRScheduler
 
+from adam_atan2_pytorch.foreach import AdamAtan2
+
 from ema_pytorch import EMA
 
 from lightning import Fabric
@@ -275,7 +277,8 @@ def __init__(
         use_ema: bool = True,
         ema_kwargs: dict = dict(
             use_foreach = True
-        )
+        ),
+        use_adam_atan2: bool = False
     ):
         super().__init__()
 
@@ -305,7 +308,13 @@ def __init__(
         # optimizer
 
         if not exists(optimizer):
-            optimizer = Adam(
+            adam_klass = Adam
+
+            if use_adam_atan2:
+                del default_adam_kwargs['eps']
+                adam_klass = AdamAtan2
+
+            optimizer = adam_klass(
                 model.parameters(),
                 lr = lr,
                 **default_adam_kwargs

pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "alphafold3-pytorch"
-version = "0.2.53"
+version = "0.2.54"
 description = "Alphafold 3 - Pytorch"
 authors = [
     { name = "Phil Wang", email = "[email protected]" }
@@ -23,6 +23,7 @@ classifiers=[
 ]
 
 dependencies = [
+    "adam-atan2-pytorch>=0.0.8",
     "beartype",
     "biopython>=1.83",
     "CoLT5-attention>=0.11.0",