Disable internal extra MSA chunking during inference

c7f26b40 · Gustaf Ahdritz · GitHub · 3279b28d · c7f26b40
Unverified Commit c7f26b40 authored Feb 03, 2022 by Gustaf Ahdritz Committed by GitHub Feb 03, 2022
Show whitespace changes
Inline Side-by-side

Showing with 1 addition and 1 deletion

openfold/model/evoformer.py openfold/model/evoformer.py +1 -1

No files found.
--- a/openfold/model/evoformer.py
+++ b/openfold/model/evoformer.py
@@ -359,7 +359,7 @@ class ExtraMSABlock(nn.Module):
                z=z.clone(), 
                mask=msa_mask, 
                chunk_size=chunk_size,
-                _chunk_logits=_chunk_logits,
+                _chunk_logits=_chunk_logits if torch.is_grad_enabled() else None,
                _checkpoint_chunks=
                    self.ckpt if torch.is_grad_enabled() else False,
            )