Hugging Face's logo

Spaces:
XAI
/
PEEB
Running on Zero

App Files Community
1
Peijie commited on
Commit
ef7d4dd
·
1 Parent(s): 11b9b46

test gpu.

Browse files
Files changed (1) hide show
  1. app.py +1 -1
app.py CHANGED
@@ -20,7 +20,7 @@ from utils.load_model import load_xclip
20
  from utils.predict import xclip_pred
21
 
22
 
23
- DEVICE = "cpu"
24
  XCLIP, OWLVIT_PRECESSOR = load_xclip(DEVICE)
25
  XCLIP_DESC_PATH = "data/jsons/bs_cub_desc.json"
26
  XCLIP_DESC = json.load(open(XCLIP_DESC_PATH, "r"))
 
20
  from utils.predict import xclip_pred
21
 
22
 
23
+ DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
24
  XCLIP, OWLVIT_PRECESSOR = load_xclip(DEVICE)
25
  XCLIP_DESC_PATH = "data/jsons/bs_cub_desc.json"
26
  XCLIP_DESC = json.load(open(XCLIP_DESC_PATH, "r"))