Spaces:

open-nlp
/

Chris-lab

Sleeping

kz209 commited on Jul 29, 2024

Commit

d5f9bcf

1 Parent(s): dea4ce7

change auto

Files changed (1) hide show

utils/model.py CHANGED Viewed

@@ -24,7 +24,7 @@ class Model(torch.nn.Module):
             model=model_name,
             tokenizer=self.tokenizer,
             torch_dtype=torch.bfloat16,
-            device_map="cpu",
         )
         self.update()

             model=model_name,
             tokenizer=self.tokenizer,
             torch_dtype=torch.bfloat16,
+            device_map="auto",
         )
         self.update()