Update modeling_markupdm.py

modeling_markupdm.py

@@ -105,12 +105,12 @@ class MarkupDMForCausalLM(PreTrainedModel, GenerationMixin):  # type: ignore
     def from_pretrained(cls, *args: Any, **kwargs: Any) -> "MarkupDMForCausalLM":
         assert "config" in kwargs, "Config must be provided"
         config = kwargs["config"]
-        torch_dtype = kwargs.get("torch_dtype", None)
+        dtype = kwargs.get("dtype", kwargs.get("torch_dtype", None))
 
         # Initialize text model
         text_model = AutoModelForCausalLM.from_config(
             config.text_model,
-            torch_dtype=torch_dtype,
+            dtype=dtype,
             attn_implementation=config._attn_implementation,
         )
 
@@ -119,7 +119,7 @@ class MarkupDMForCausalLM(PreTrainedModel, GenerationMixin):  # type: ignore
             vision_model = AutoModel.from_config(
                 config.vision_model,
                 trust_remote_code=True,
-                torch_dtype=torch_dtype,
+                dtype=dtype,
             )
 
         return super().from_pretrained(  # type: ignore