Update ultravox_model.py

ultravox_model.py

@@ -76,6 +76,9 @@ class UltravoxModel(transformers.LlamaPreTrainedModel, GenerationMixin):
         return model
 
     def _load_child_model_weights(self, *args, **kwargs) -> "UltravoxModel":
+        if "torch_dtype" in kwargs:
+            self.config.torch_dtype = kwargs.pop("torch_dtype")
+
         if (
             self.config.text_model_id is not None
             and self.language_model.device.type == "meta"
@@ -850,4 +853,4 @@ UltravoxModel.register_for_auto_class()
 transformers.AutoConfig.register("ultravox", UltravoxConfig)
 transformers.AutoModel.register(UltravoxConfig, UltravoxModel)
 
-transformers.activations.ACT2FN["swiglu"] = SwiGLU
+transformers.activations.ACT2FN["swiglu"] = SwiGLU