new

models/llm/llm.py

@@ -42,9 +42,9 @@ OUTPUT ENGLISH ONLY! NOT ANY OTHER LANGUAGE, E.G., CHINESE!"""
 def load_llm_model(model_name_or_dir, torch_dtype='auto', device_map='cpu'):
     model = AutoModelForCausalLM.from_pretrained(
         model_name_or_dir,
-        torch_dtype=torch_dtype,
+        # torch_dtype=torch_dtype,
         # torch_dtype=torch.float8_e5m2,
-        # torch_dtype=torch.float16,
+        torch_dtype=torch.float16,
         device_map=device_map
     )
     print(f'set llm model to {model_name_or_dir}')