HF_Final_Assignment_Template

Sleeping

Samuel Thomas commited on May 31

Commit

8792b9a

1 Parent(s): fe1bd6e

cpu:

Files changed (1) hide show

tools.py CHANGED Viewed

@@ -102,7 +102,7 @@ def create_llm_pipeline():
     return pipeline(
         "text-generation",
         model=model_id,
-        device_map="auto",
         torch_dtype=torch.float16,
         max_new_tokens=1024,
         temperature=0.05,

     return pipeline(
         "text-generation",
         model=model_id,
+        device_map="cpu",
         torch_dtype=torch.float16,
         max_new_tokens=1024,
         temperature=0.05,