abenkbp commited on
Commit
b0b3de5
·
1 Parent(s): 5180d45
Files changed (1) hide show
  1. data/models/llama3-1-8b.py +1 -1
data/models/llama3-1-8b.py CHANGED
@@ -16,7 +16,7 @@ pipeline = transformers.pipeline(
16
  "text-generation",
17
  model=model_id,
18
  model_kwargs={"torch_dtype": torch.bfloat16},
19
- device="gpu",
20
  token=api_key
21
  )
22
 
 
16
  "text-generation",
17
  model=model_id,
18
  model_kwargs={"torch_dtype": torch.bfloat16},
19
+ device="cuda",
20
  token=api_key
21
  )
22