Disable activation checkpointing during inference

openfold/model/evoformer.py

@@ -350,7 +350,7 @@ class EvoformerStack(nn.Module):
                 ) for b in self.blocks
             ], 
             args=(m, z),
-            blocks_per_ckpt=self.blocks_per_ckpt,
+            blocks_per_ckpt=self.blocks_per_ckpt if self.training else None,
         )
 
         s = None

openfold/model/template.py

@@ -265,7 +265,7 @@ class TemplatePairStack(nn.Module):
                 ) for b in self.blocks
             ], 
             args=(t),
-            blocks_per_ckpt=self.blocks_per_ckpt,
+            blocks_per_ckpt=self.blocks_per_ckpt if self.training else None,
         )
 
         t = self.layer_norm(t)