HF_Final_Assignment_Template

Sleeping

Samuel Thomas commited on May 25

Commit

f2a058a

1 Parent(s): 5948f01

move to cpu

Files changed (1) hide show

tools.py CHANGED Viewed

@@ -83,7 +83,7 @@ def create_llm_pipeline():
         "text-generation",
         model=model_id,
         tokenizer = tokenizer,
-        device_map="auto",
         torch_dtype=torch.float16,
         max_new_tokens=1024,
         temperature=0.1

         "text-generation",
         model=model_id,
         tokenizer = tokenizer,
+        device_map="cpu",
         torch_dtype=torch.float16,
         max_new_tokens=1024,
         temperature=0.1