is_scripting() guard on checkpoint_seq

2 years ago · 72fba669a8
parent 95ec255f7f
commit 72fba669a8
2 changed files with 2 additions and 2 deletions
--- a/timm/models/efficientformer.py
+++ b/timm/models/efficientformer.py
@ -336,7 +336,7 @@ class EfficientFormerStage(nn.Module):
    def forward(self, x):
        x = self.downsample(x)
-        if self.grad_checkpointing:
+        if self.grad_checkpointing and not torch.jit.is_scripting():
            x = checkpoint_seq(self.blocks, x)
        else:
            x = self.blocks(x)
--- a/timm/models/efficientformer_v2.py
+++ b/timm/models/efficientformer_v2.py
@ -499,7 +499,7 @@ class EfficientFormerV2Stage(nn.Module):
    def forward(self, x):
        x = self.downsample(x)
-        if self.grad_checkpointing:
+        if self.grad_checkpointing and not torch.jit.is_scripting():
            x = checkpoint_seq(self.blocks, x)
        else:
            x = self.blocks(x)