bypass cuda error
Browse files
app.py
CHANGED
|
@@ -61,7 +61,10 @@ def get_pipeline(model_name):
|
|
| 61 |
mdl = AutoModelForCausalLM.from_pretrained(
|
| 62 |
model_name, weights_only=False, trust_remote_code=True
|
| 63 |
)
|
| 64 |
-
|
|
|
|
|
|
|
|
|
|
| 65 |
return pipeline("text-generation", model=mdl, tokenizer=tok, device=0)
|
| 66 |
|
| 67 |
@spaces.GPU
|
|
|
|
| 61 |
mdl = AutoModelForCausalLM.from_pretrained(
|
| 62 |
model_name, weights_only=False, trust_remote_code=True
|
| 63 |
)
|
| 64 |
+
try:
|
| 65 |
+
mdl.to("cuda")
|
| 66 |
+
except Exception as e:
|
| 67 |
+
print(f'Error: {e}')
|
| 68 |
return pipeline("text-generation", model=mdl, tokenizer=tok, device=0)
|
| 69 |
|
| 70 |
@spaces.GPU
|