xinyun/ComfyUI
1
0
Fork 0
mirror of https://git.datalinker.icu/comfyanonymous/ComfyUI synced 2026-09-03 21:17:08 +08:00
Code Issues Packages Projects Releases Wiki Activity

Remove grad ckpt from model impl

Browse Source 
we apply grad ckpt in train node directly
This commit is contained in:
Kohaku-Blueleaf 2025-06-13 17:51:37 +08:00 committed by GitHub
parent 31c8cc9aa3
commit bbcc65e7e0
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
1 changed files with 8 additions and 11 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

19
comfy/ldm/cosmos/blocks.py
View File

@ -796,15 +796,12 @@ class GeneralDITTransformerBlock(nn.Module):
adaln_lora_B_3D: Optional[torch.Tensor] = None,
) -> torch.Tensor:
for block in self.blocks:
if self.training:
x = torch.utils.checkpoint.checkpoint(block, x, emb_B_D, crossattn_emb, crossattn_mask, rope_emb_L_1_1_D, adaln_lora_B_3D, use_reentrant=False)
else:
x = block(
x,
emb_B_D,
crossattn_emb,
crossattn_mask,
rope_emb_L_1_1_D=rope_emb_L_1_1_D,
adaln_lora_B_3D=adaln_lora_B_3D,
)
x = block(
x,
emb_B_D,
crossattn_emb,
crossattn_mask,
rope_emb_L_1_1_D=rope_emb_L_1_1_D,
adaln_lora_B_3D=adaln_lora_B_3D,
)
return x