Update app.py
Browse files
app.py
CHANGED
|
@@ -37,7 +37,7 @@ def prompt_model(model_id, system_prompt, user_prompt):
|
|
| 37 |
{"role": "assistant", "content": ""}
|
| 38 |
]
|
| 39 |
|
| 40 |
-
output = pipe(messages, device="cuda")
|
| 41 |
|
| 42 |
return output[0]["generated_text"][-1]["content"]
|
| 43 |
|
|
|
|
| 37 |
{"role": "assistant", "content": ""}
|
| 38 |
]
|
| 39 |
|
| 40 |
+
output = pipe(messages, model_kwargs={"torch_dtype": torch.bfloat16}, device="cuda")
|
| 41 |
|
| 42 |
return output[0]["generated_text"][-1]["content"]
|
| 43 |
|