Chinese-Llama-2-7b

Runtime error

shiyemin2 commited on Jul 22, 2023

Commit

fa0f42c

•

1 Parent(s): cef1f39

Update model.py

Files changed (1) hide show

model.py CHANGED Viewed

@@ -14,7 +14,8 @@ if torch.cuda.is_available():
         model = AutoModelForCausalLM.from_pretrained(
                 model_id,
                 load_in_4bit=True,
-                torch_dtype=torch.float16
             )
     else:
         model = AutoModelForCausalLM.from_pretrained(

         model = AutoModelForCausalLM.from_pretrained(
                 model_id,
                 load_in_4bit=True,
+                torch_dtype=torch.float16,
+                device_map='auto'
             )
     else:
         model = AutoModelForCausalLM.from_pretrained(