More forward_intermediate specific grad checkpointing fixes

Author: rwightman

timm/models/starnet.py

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

timm/models/tiny_vit.py

@@ -571,7 +571,7 @@ def forward_intermediates(
 
         for feat_idx, stage in enumerate(stages):
             if self.grad_checkpointing and not torch.jit.is_scripting():
-                x = checkpoint(stages, x)
+                x = checkpoint(stage, x)
             else:
                 x = stage(x)
             if feat_idx in take_indices: