Spaces:
Sleeping
Sleeping
test gpu.
Browse files
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"))
|