You need to sign in or sign up before continuing.
Unverified Commit c7f26b40 authored by Gustaf Ahdritz's avatar Gustaf Ahdritz Committed by GitHub
Browse files

Disable internal extra MSA chunking during inference

parent 3279b28d
...@@ -359,7 +359,7 @@ class ExtraMSABlock(nn.Module): ...@@ -359,7 +359,7 @@ class ExtraMSABlock(nn.Module):
z=z.clone(), z=z.clone(),
mask=msa_mask, mask=msa_mask,
chunk_size=chunk_size, chunk_size=chunk_size,
_chunk_logits=_chunk_logits, _chunk_logits=_chunk_logits if torch.is_grad_enabled() else None,
_checkpoint_chunks= _checkpoint_chunks=
self.ckpt if torch.is_grad_enabled() else False, self.ckpt if torch.is_grad_enabled() else False,
) )
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment