kiri00 commited on
Commit
5f76d85
1 Parent(s): d3f1dd3

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +1 -1
app.py CHANGED
@@ -36,7 +36,7 @@ def predict(image, max_length=64, num_beams=4):
36
  return text
37
 
38
  model_path = "team-indain-image-caption/hindi-image-captioning"
39
- device = torch.device("cpu")
40
  # Load model.
41
  model = VisionEncoderDecoderModel.from_pretrained(model_path)
42
  model.to(device)
 
36
  return text
37
 
38
  model_path = "team-indain-image-caption/hindi-image-captioning"
39
+ device = torch.device("cuda:0" if torch.cuda.is_available() else"cpu")
40
  # Load model.
41
  model = VisionEncoderDecoderModel.from_pretrained(model_path)
42
  model.to(device)