Change Emu2 weight loading to use proper torch.bfloat16

kushal-10 · Apr 30, 2024 · d191b80 · d191b80
1 parent ff86080
commit d191b80
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/backends/huggingface_multimodal_api.py b/backends/huggingface_multimodal_api.py
@@ -64,7 +64,7 @@ def load_model(model_spec: backends.ModelSpec):
     if hasattr(model_spec, 'trust_remote_code'):
         if model_spec['trust_remote_code']:
             if model_spec['model_type'] == "Emu2":
-                model = model_type.from_pretrained(hf_model_str, device_map="auto", torch_dtype="bfloat16",
+                model = model_type.from_pretrained(hf_model_str, device_map="auto", torch_dtype=torch.bfloat16,
                                                    trust_remote_code=model_spec['trust_remote_code'])
             else:
                 model = model_type.from_pretrained(hf_model_str, device_map="auto", torch_dtype="auto",