remove compilation

Author: ardagoreci

src/models/diffusion_module.py

@@ -158,7 +158,6 @@ 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)

src/models/model.py

@@ -59,7 +59,6 @@ def __init__(self, config):
             LinearNoBias(self.c_z, self.c_z)
         )
 
-    @torch.compile
     def run_trunk(
             self,
             feats: Dict[str, Tensor],