Spaces:

xset
/

aibot

Sleeping

xset commited on May 18

Commit

e4fa82c

1 Parent(s): 85e9a75

float16

Files changed (1) hide show

main.py CHANGED Viewed

@@ -23,7 +23,7 @@ try:
     tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
     model = AutoModelForCausalLM.from_pretrained(
         MODEL_NAME,
-        torch_dtype=torch.float32,
         device_map="auto",
         low_cpu_mem_usage=True
     )

     tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
     model = AutoModelForCausalLM.from_pretrained(
         MODEL_NAME,
+        torch_dtype=torch.float16,
         device_map="auto",
         low_cpu_mem_usage=True
     )