support Anima gradient checkpointing (#1319)

This commit is contained in:
Zhongjie Duan
2026-03-02 19:06:55 +08:00
committed by GitHub
parent f43b18ec21
commit b3ef224042

View File

@@ -256,6 +256,8 @@ def model_fn_anima(
timesteps=timestep,
context=prompt_emb,
t5xxl_ids=t5xxl_ids,
use_gradient_checkpointing=use_gradient_checkpointing,
use_gradient_checkpointing_offload=use_gradient_checkpointing_offload,
)
model_output = model_output.squeeze(2)
return model_output