From bbdb05d744b0b3f78dc0c8de6f56bfc61ad23145 Mon Sep 17 00:00:00 2001
From: Jaret Burkett <jaretburkett@gmail.com>
Date: Sun, 10 May 2026 12:10:02 -0600
Subject: [PATCH] Fix saving for hidream-o1

---
 .../hidream/src/hidream_o1/qwen3_vl_transformers.py             | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

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 499a6de0..0ec7695b 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
@@ -2104,7 +2104,7 @@ class Qwen3VLCausalLMOutputWithPast(ModelOutput):
 
 class Qwen3VLForConditionalGeneration(Qwen3VLPreTrainedModel, GenerationMixin):
     _checkpoint_conversion_mapping = {}
-    _tied_weights_keys = ["lm_head.weight"]
+    _tied_weights_keys = {"lm_head.weight": "model.language_model.embed_tokens.weight"}
     # Reference: fix gemma3 grad acc #37208
     accepts_loss_kwargs = False
     config: Qwen3VLConfig