switch to dynamic gpu
Browse files
app.py
CHANGED
|
@@ -27,7 +27,7 @@ MODEL_STANFORD_LOCATION = hf_hub_download(
|
|
| 27 |
filename="checkpoint_best_stanford.pth"
|
| 28 |
)
|
| 29 |
|
| 30 |
-
DEVICE = "cpu"
|
| 31 |
|
| 32 |
DATASET_CHECKPOINT_MAP = {
|
| 33 |
"Flowers": MODEL_FLOWERS_LOCATION,
|
|
|
|
| 27 |
filename="checkpoint_best_stanford.pth"
|
| 28 |
)
|
| 29 |
|
| 30 |
+
DEVICE = torch.device("cuda") if torch.cuda.is_available() else torch.device("cpu")
|
| 31 |
|
| 32 |
DATASET_CHECKPOINT_MAP = {
|
| 33 |
"Flowers": MODEL_FLOWERS_LOCATION,
|