add compilation

Ligo-Biosciences · Aug 24, 2024 · fa95f99 · fa95f99
1 parent be9acdf
commit fa95f99
Show file tree

Hide file tree

Showing 3 changed files with 2 additions and 176 deletions.
diff --git a/configs/model/mini-af3.yaml b/configs/model/mini-af3.yaml
diff --git a/src/models/diffusion_module.py b/src/models/diffusion_module.py
@@ -158,6 +158,7 @@ def rescale_with_updates(
         r_update_scale = torch.sqrt(noisy_pos_scale) * timesteps.unsqueeze(-1)
         return noisy_atoms * noisy_pos_scale + r_updates * r_update_scale
 
+    @torch.compile
     def forward(
             self,
             noisy_atoms: Tensor,  # (bs, S, n_atoms, 3)

diff --git a/src/models/model.py b/src/models/model.py
@@ -59,6 +59,7 @@ def __init__(self, config):
             LinearNoBias(self.c_z, self.c_z)
         )
 
+    @torch.compile
     def run_trunk(
             self,
             feats: Dict[str, Tensor],