xinyun/ComfyUI
1
0
Fork 0
mirror of https://git.datalinker.icu/comfyanonymous/ComfyUI synced 2026-09-04 02:17:05 +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

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

@ -796,9 +796,6 @@ class GeneralDITTransformerBlock(nn.Module):
adaln_lora_B_3D: Optional[torch.Tensor] = None, adaln_lora_B_3D: Optional[torch.Tensor] = None,
) -> torch.Tensor: ) -> torch.Tensor:
for block in self.blocks: 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 = block(
x, x,
emb_B_D, emb_B_D,