Update app.py

app.py

@@ -73,7 +73,7 @@ feature_extractor = ViTImageProcessor.from_pretrained(encoder_name_or_path)
 
 
 
-# device = "cuda" if torch.cuda.is_available() else "cpu"
+device = "cuda" if torch.cuda.is_available() else "cpu"
 
 # cap_model.to(device)
 
@@ -101,7 +101,7 @@ def predict_event(image):
 
     # img =  Image.open(image).convert("RGB")
 
-    generated_caption = tokenizer.decode(model.generate(feature_extractor(image, return_tensors="pt").pixel_values.to("cuda"))[0])
+    generated_caption = tokenizer.decode(model.generate(feature_extractor(image, return_tensors="pt").pixel_values.to(device))[0])
 
     # caption_vitgpt = generate_caption(model, image)
     #caption_vitgpt = generate_caption(feature_extractor, cap_model, image, tokenizer)