Spaces:

factorstudios
/

NEWM

Runtime error

Factor Studios commited on Aug 6

Commit

59a2bc8

verified ·

1 Parent(s): 3679667

Update src/routes/huggingface_gpt_model.py

Files changed (1) hide show

src/routes/huggingface_gpt_model.py CHANGED Viewed

@@ -35,7 +35,7 @@ class HuggingFaceGPTModel:
             self.model = AutoModelForCausalLM.from_pretrained(
                 model_name,
                 torch_dtype=torch.float32,
-                device_map="cpu",
                 low_cpu_mem_usage=True
             )

             self.model = AutoModelForCausalLM.from_pretrained(
                 model_name,
                 torch_dtype=torch.float32,
+                device_map="gpu",
                 low_cpu_mem_usage=True
             )