Fix several grad checkpointing issues

Author: rwightman

timm/models/efficientnet.py

@@ -261,7 +261,7 @@ def forward_intermediates(
             blocks = self.blocks[:max_index]
         for feat_idx, blk in enumerate(blocks, start=1):
             if self.grad_checkpointing and not torch.jit.is_scripting():
-                x = checkpoint_seq(blk, x, flatten=True)
+                x = checkpoint_seq(blk, x)
             else:
                 x = blk(x)
             if feat_idx in take_indices:

timm/models/fasternet.py

@@ -142,7 +142,7 @@ def __init__(
     def forward(self, x: torch.Tensor) -> torch.Tensor:
         x = self.downsample(x)
         if self.grad_checkpointing and not torch.jit.is_scripting():
-            x = checkpoint_seq(self.blocks, x, flatten=True)
+            x = checkpoint_seq(self.blocks, x)
         else:
             x = self.blocks(x)
         return x

timm/models/focalnet.py

@@ -274,7 +274,7 @@ def forward(self, x):
         x = self.downsample(x)
         for blk in self.blocks:
             if self.grad_checkpointing and not torch.jit.is_scripting():
-                x = checkpoint.checkpoint(blk, x)
+                x = checkpoint(blk, x)
             else:
                 x = blk(x)
         return x

timm/models/gcvit.py

@@ -361,7 +361,7 @@ def forward(self, x):
         global_query = self.global_norm(global_query.permute(0, 2, 3, 1))
         for blk in self.blocks:
             if self.grad_checkpointing and not torch.jit.is_scripting():
-                x = checkpoint.checkpoint(blk, x)
+                x = checkpoint(blk, x, global_query)
             else:
                 x = blk(x, global_query)
         x = self.norm(x)

timm/models/ghostnet.py

@@ -728,7 +728,7 @@ def forward_intermediates(
 
         for feat_idx, stage in enumerate(stages, start=1):
             if self.grad_checkpointing and not torch.jit.is_scripting():
-                x = checkpoint_seq(stage, x, flatten=True)
+                x = checkpoint_seq(stage, x)
             else:
                 x = stage(x)
             if feat_idx in take_indices:

timm/models/hgnet.py

@@ -345,7 +345,7 @@ def __init__(
     def forward(self, x):
         x = self.downsample(x)
         if self.grad_checkpointing and not torch.jit.is_scripting():
-            x = checkpoint_seq(self.blocks, x, flatten=False)
+            x = checkpoint_seq(self.blocks, x)
         else:
             x = self.blocks(x)
         return x

timm/models/hieradet_sam2.py

@@ -288,6 +288,7 @@ def __init__(
         norm_layer = get_norm_layer(norm_layer)
         act_layer = get_act_layer(act_layer)
         assert len(stages) == len(window_spec)
+        self.grad_checkpointing = False
         self.num_classes = num_classes
         self.window_spec = window_spec
         self.output_fmt = 'NHWC'

timm/models/mobilenetv3.py

@@ -229,7 +229,7 @@ def forward_intermediates(
             blocks = self.blocks[:max_index]
         for feat_idx, blk in enumerate(blocks, start=1):
             if self.grad_checkpointing and not torch.jit.is_scripting():
-                x = checkpoint_seq(blk, x, flatten=True)
+                x = checkpoint_seq(blk, x)
             else:
                 x = blk(x)
             if feat_idx in take_indices:

timm/models/mvitv2.py

@@ -681,7 +681,7 @@ def __init__(
     def forward(self, x, feat_size: List[int]):
         for blk in self.blocks:
             if self.grad_checkpointing and not torch.jit.is_scripting():
-                x, feat_size = checkpoint.checkpoint(blk, x, feat_size)
+                x, feat_size = checkpoint(blk, x, feat_size)
             else:
                 x, feat_size = blk(x, feat_size)
         return x, feat_size

timm/models/pvt_v2.py

@@ -267,7 +267,7 @@ def forward(self, x):
         x = x.reshape(B, -1, C)
         for blk in self.blocks:
             if self.grad_checkpointing and not torch.jit.is_scripting():
-                x = checkpoint.checkpoint(blk, x, feat_size)
+                x = checkpoint(blk, x, feat_size)
             else:
                 x = blk(x, feat_size)
         x = self.norm(x)