From af82197907a9dc2a11a56961e0483157fd3ed652 Mon Sep 17 00:00:00 2001 From: Jaret Burkett Date: Tue, 8 Sep 2026 12:24:56 -0600 Subject: [PATCH] Fix gradient checkpointing for sapiens DFE --- toolkit/models/diffusion_feature_extraction.py | 1 + toolkit/models/sapiens2.py | 3 ++- 2 files changed, 3 insertions(+), 1 deletion(-) diff --git a/toolkit/models/diffusion_feature_extraction.py b/toolkit/models/diffusion_feature_extraction.py index b1ba570..2f6ced6 100644 --- a/toolkit/models/diffusion_feature_extraction.py +++ b/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 diff --git a/toolkit/models/sapiens2.py b/toolkit/models/sapiens2.py index 9f4fb28..90d990c 100644 --- a/toolkit/models/sapiens2.py +++ b/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)