hanzla commited on
Commit
851843d
·
1 Parent(s): 77ab3f6
Files changed (1) hide show
  1. src/pdfchatbot.py +1 -1
src/pdfchatbot.py CHANGED
@@ -59,7 +59,7 @@ class PDFChatBot:
59
  "text-generation",
60
  model="meta-llama/Meta-Llama-3-8B-Instruct",
61
  model_kwargs={"torch_dtype": torch.float16},
62
- device="cuda",
63
  )
64
  print("Model pipeline loaded")
65
 
 
59
  "text-generation",
60
  model="meta-llama/Meta-Llama-3-8B-Instruct",
61
  model_kwargs={"torch_dtype": torch.float16},
62
+ device_map="auto",
63
  )
64
  print("Model pipeline loaded")
65