enabled auto torchdtype detecting

Signed-off-by: ftgreat <[email protected]>
capjamesg · Oct 9, 2023 · 3c0cb59 · 3c0cb59
1 parent df9c0c2
commit 3c0cb59
Showing 1 changed file with 7 additions and 1 deletion.
diff --git a/flagai/auto_model/auto_loader.py b/flagai/auto_model/auto_loader.py
@@ -210,7 +210,13 @@ def __init__(self,
         if task_name == "aquila2":
             from flagai.model.aquila2.modeling_aquila import AquilaForCausalLM
             download_path = os.path.join(model_dir, model_name)
-
+
+            if not torch_dtype:  
+                if model_name.lower() == "aquilachat2-34b":
+                    torch_dtype = torch.bfloat16
+                else:
+                    torch_dtype = torch.float16
+
             if not os.path.exists(download_path):
                 # Try to download from ModelHub
                 try: