Fix gradient checkpointing for sapiens DFE

toolkit/models/diffusion_feature_extraction.py

@@ -1094,6 +1094,7 @@ class DiffusionFeatureExtractor9(nn.Module):
         self.model = Sapiens2(arch="sapiens2_1b", img_size=(1024, 768), patch_size=16).eval().cuda()  # img_size is (H, W)
         self.model.load_state_dict(load_file(ckpt_path))
         self.model.to(device, dtype=dtype)
+        self.model.enable_gradient_checkpointing()
 
         self.losses = {}
         self.log_every = 100

toolkit/models/sapiens2.py

@@ -766,7 +766,8 @@ class Sapiens2(nn.Module):
         rope_sincos = self.rope_embed(H=patch_resolution[0], W=patch_resolution[1])
         outs = []
         for i, layer in enumerate(self.blocks):
-            if self.gradient_checkpointing and self.training:
+            # gate on is_grad_enabled, not self.training: DFE keeps this in eval
+            if self.gradient_checkpointing and torch.is_grad_enabled():
                 x = checkpoint(layer, x, rope_sincos, use_reentrant=False)
             else:
                 x = layer(x, rope=rope_sincos)