fix metaformer and nextvit

Author: brianhou0208

timm/models/metaformer.py

@@ -41,7 +41,7 @@
     use_fused_attn
 from ._builder import build_model_with_cfg
 from ._features import feature_take_indices
-from ._manipulate import checkpoint_seq
+from ._manipulate import checkpoint, checkpoint_seq
 from ._registry import generate_default_cfgs, register_model
 
 __all__ = ['MetaFormer']
@@ -631,8 +631,8 @@ def forward_intermediates(
             stages = self.stages[:max_index + 1]
 
         for feat_idx, stage in enumerate(stages):
-            if self.grad_checkpointing and stage.grad_checkpointing and not torch.jit.is_scripting():
-                x = checkpoint_seq(stage, x)
+            if self.grad_checkpointing and not torch.jit.is_scripting():
+                x = checkpoint(stage, x)
             else:
                 x = stage(x)
             if feat_idx in take_indices:

timm/models/nextvit.py

@@ -17,7 +17,7 @@
 from timm.layers import ClassifierHead
 from ._builder import build_model_with_cfg
 from ._features import feature_take_indices
-from ._manipulate import checkpoint_seq
+from ._manipulate import checkpoint, checkpoint_seq
 from ._registry import generate_default_cfgs, register_model
 
 __all__ = ['NextViT']
@@ -594,7 +594,10 @@ def forward_intermediates(
             stages = self.stages[:max_index + 1]
 
         for feat_idx, stage in enumerate(stages):
-            x = stage(x)
+            if self.grad_checkpointing and not torch.jit.is_scripting():
+                x = checkpoint(stage, x)
+            else:
+                x = stage(x)
             if feat_idx in take_indices:
                 if feat_idx == last_idx:
                     x_inter = self.norm(x) if norm else x