NaFlexVit and NextVit forward_intermediates grad checkpointing fixes

Author: rwightman

timm/models/naflexvit.py

@@ -1257,7 +1257,7 @@ def forward_intermediates(
             if attn_mask is not None:
                 x = blk(x, attn_mask=attn_mask)
             elif self.grad_checkpointing and not torch.jit.is_scripting():
-                x = checkpoint(blk. x)
+                x = checkpoint(blk, x)
             else:
                 x = blk(x)
             if i 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']
@@ -595,7 +595,7 @@ def forward_intermediates(
 
         for feat_idx, stage in enumerate(stages):
             if self.grad_checkpointing and not torch.jit.is_scripting():
-                x = checkpoint_seq(stage, x)
+                x = checkpoint(stage, x)
             else:
                 x = stage(x)
             if feat_idx in take_indices: