Spaces:

orionweller
/

rank1-test-time-compute

Running on Zero

orionweller commited on 5 days ago

Commit

776166c

verified ·

1 Parent(s): b5f72c9

Update model.py

Files changed (1) hide show

model.py CHANGED Viewed

@@ -52,7 +52,7 @@ class Rank1:
         # Load AWQ model on CPU initially
         self.model = AutoModelForCausalLM.from_pretrained(
             model_name_or_path,
-            device_map="cpu",
             trust_remote_code=True,
             attn_implementation="flash_attention_2"
         )

         # Load AWQ model on CPU initially
         self.model = AutoModelForCausalLM.from_pretrained(
             model_name_or_path,
+            device_map="auto",
             trust_remote_code=True,
             attn_implementation="flash_attention_2"
         )