mirrors/pytorch-image-models
1
0
Fork 0
mirror of https://github.com/huggingface/pytorch-image-models.git synced 2025-06-03 15:01:08 +08:00
Code Issues Projects Releases Wiki Activity

is_scripting() guard on checkpoint_seq

Browse Source
This commit is contained in:
Ross Wightman 2023-02-04 14:21:49 -08:00
parent 95ec255f7f
commit 72fba669a8
2 changed files with 2 additions and 2 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

2
timm/models/efficientformer.py
View File

@ -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)

2
timm/models/efficientformer_v2.py
View File

@ -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)