Spaces:
Sleeping
Sleeping
Commit
·
e993e1b
1
Parent(s):
7e9dee7
fix: pipeline does not have 'to'
Browse files- prediction.py +1 -1
prediction.py
CHANGED
|
@@ -28,8 +28,8 @@ pipeline = hf_pipeline(
|
|
| 28 |
"text-generation",
|
| 29 |
model="meta-llama/Meta-Llama-3.1-8B-Instruct",
|
| 30 |
model_kwargs={"torch_dtype": 'bfloat16'},
|
|
|
|
| 31 |
)
|
| 32 |
-
pipeline.to('cuda')
|
| 33 |
|
| 34 |
|
| 35 |
class ModelPrediction:
|
|
|
|
| 28 |
"text-generation",
|
| 29 |
model="meta-llama/Meta-Llama-3.1-8B-Instruct",
|
| 30 |
model_kwargs={"torch_dtype": 'bfloat16'},
|
| 31 |
+
device_map="auto",
|
| 32 |
)
|
|
|
|
| 33 |
|
| 34 |
|
| 35 |
class ModelPrediction:
|