Spaces:
Build error
Build error
Celine commited on
Commit ·
f290e10
1
Parent(s): e5af996
comment out to.cuda()
Browse files
app.py
CHANGED
|
@@ -18,7 +18,7 @@ gemma_model = AutoModelForCausalLM.from_pretrained(
|
|
| 18 |
|
| 19 |
@spaces.GPU
|
| 20 |
def promptGemma2(textprompt, max_tokens=20):
|
| 21 |
-
input_ids = gemma_tokenizer(textprompt, return_tensors="pt").to("cuda")
|
| 22 |
outputs = gemma_model.generate(**input_ids, max_new_tokens=max_tokens)
|
| 23 |
return gemma_tokenizer.decode(outputs[0])
|
| 24 |
|
|
|
|
| 18 |
|
| 19 |
@spaces.GPU
|
| 20 |
def promptGemma2(textprompt, max_tokens=20):
|
| 21 |
+
input_ids = gemma_tokenizer(textprompt, return_tensors="pt")#.to("cuda")
|
| 22 |
outputs = gemma_model.generate(**input_ids, max_new_tokens=max_tokens)
|
| 23 |
return gemma_tokenizer.decode(outputs[0])
|
| 24 |
|