chats-bug commited on
Commit
7d06c4c
1 Parent(s): bc65b96

Updated device type

Browse files
Files changed (1) hide show
  1. app.py +2 -1
app.py CHANGED
@@ -1,4 +1,5 @@
1
  import gradio as gr
 
2
  from PIL import Image
3
 
4
  from model import GitBaseCocoModel
@@ -26,7 +27,7 @@ def generate_captions(
26
  list[str]
27
  """
28
 
29
- device = "cuda" if gradio.use_gpu else "cpu"
30
  checkpoint = "microsoft/git-base-coco"
31
 
32
  model = GitBaseCocoModel(device, checkpoint)
 
1
  import gradio as gr
2
+ import torch
3
  from PIL import Image
4
 
5
  from model import GitBaseCocoModel
 
27
  list[str]
28
  """
29
 
30
+ device = "cuda" if torch.cuda.is_available() else "cpu"
31
  checkpoint = "microsoft/git-base-coco"
32
 
33
  model = GitBaseCocoModel(device, checkpoint)