aavetis commited on
Commit
6b230a2
1 Parent(s): b9c1647

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +1 -1
app.py CHANGED
@@ -22,7 +22,7 @@ def generate_caption(image, prompt):
22
  max_new_tokens=128,
23
  eos_token_id=32001,
24
  pad_token_id=processor.tokenizer.pad_token_id
25
- )
26
 
27
  prompt_len = inputs["input_ids"].shape[1]
28
  decoded_text = processor.batch_decode(output[:, prompt_len:])[0]
 
22
  max_new_tokens=128,
23
  eos_token_id=32001,
24
  pad_token_id=processor.tokenizer.pad_token_id
25
+ ).to('cuda')
26
 
27
  prompt_len = inputs["input_ids"].shape[1]
28
  decoded_text = processor.batch_decode(output[:, prompt_len:])[0]