detach when recycling by default

Author: lucidrains

alphafold3_pytorch/alphafold3.py

@@ -4518,6 +4518,7 @@ def __init__(
         checkpoint_input_embedding = False,
         checkpoint_trunk_pairformer = False,
         checkpoint_diffusion_token_transformer = False,
+        detach_when_recycling = True
     ):
         super().__init__()
 
@@ -4642,6 +4643,8 @@ def __init__(
 
         # recycling related
 
+        self.detach_when_recycling = detach_when_recycling
+
         self.recycle_single = nn.Sequential(
             nn.LayerNorm(dim_single),
             LinearNoBias(dim_single, dim_single)
@@ -4835,7 +4838,8 @@ def forward(
         return_present_sampled_atoms: bool = False,
         return_confidence_head_logits: bool = False,
         num_rollout_steps: int | None = None,
-        rollout_show_tqdm_pbar: bool = False
+        rollout_show_tqdm_pbar: bool = False,
+        detach_when_recycling: bool = None
     ) -> (
         Float['b m 3'] |
         Float['l 3'] |
@@ -4995,6 +4999,9 @@ def forward(
 
         # init recycled single and pairwise
 
+        detach_when_recycling = default(detach_when_recycling, self.detach_when_recycling)
+        maybe_recycling_detach = torch.detach if detach_when_recycling else identity
+
         recycled_pairwise = recycled_single = None
         single = pairwise = None
 
@@ -5007,9 +5014,11 @@ def forward(
             recycled_single = recycled_pairwise = 0.
 
             if exists(single):
+                single = maybe_recycling_detach(single)
                 recycled_single = self.recycle_single(single)
 
             if exists(pairwise):
+                pairwise = maybe_recycling_detach(pairwise)
                 recycled_pairwise = self.recycle_pairwise(pairwise)
 
             single = single_init + recycled_single

pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "alphafold3-pytorch"
-version = "0.2.106"
+version = "0.2.107"
 description = "Alphafold 3 - Pytorch"
 authors = [
     { name = "Phil Wang", email = "[email protected]" },