diff --git a/extensions_built_in/diffusion_models/omnigen2/__init__.py b/extensions_built_in/diffusion_models/omnigen2/__init__.py
index ee37a1f0..ddff811e 100644
--- a/extensions_built_in/diffusion_models/omnigen2/__init__.py
+++ b/extensions_built_in/diffusion_models/omnigen2/__init__.py
@@ -84,7 +84,7 @@ class OmniGen2Model(BaseModel):
             subfolder="mllm",
             torch_dtype=torch.bfloat16
         )
-        
+        mllm.to(self.device_torch, dtype=dtype)
         if self.model_config.quantize_te:
             self.print_and_status_update("Quantizing Qwen2.5 VL model")
             quantization_type = get_qtype(self.model_config.qtype_te)