From a8001e9a3f04b3faa6a6dabb0bb127e2ba29580b Mon Sep 17 00:00:00 2001 From: Jaret Burkett Date: Sun, 10 May 2026 14:41:31 -0600 Subject: [PATCH] Fix gradient checkpointing for hidream o1 --- .../hidream/src/hidream_o1/qwen3_vl_transformers.py | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/extensions_built_in/diffusion_models/hidream/src/hidream_o1/qwen3_vl_transformers.py b/extensions_built_in/diffusion_models/hidream/src/hidream_o1/qwen3_vl_transformers.py index 26128e37..0551aeaa 100644 --- a/extensions_built_in/diffusion_models/hidream/src/hidream_o1/qwen3_vl_transformers.py +++ b/extensions_built_in/diffusion_models/hidream/src/hidream_o1/qwen3_vl_transformers.py @@ -2148,6 +2148,10 @@ class Qwen3VLForConditionalGeneration(Qwen3VLPreTrainedModel, GenerationMixin): def visual(self): return self.model.visual + def enable_gradient_checkpointing(self): + self.model.visual.gradient_checkpointing_enable() + self.model.language_model.gradient_checkpointing_enable() + @check_model_inputs def forward( self,