Spaces:
Sleeping
Sleeping
fix gpu
Browse files
app.py
CHANGED
|
@@ -17,7 +17,6 @@ from PIL import Image
|
|
| 17 |
args = DotMap()
|
| 18 |
args.deploy = 'vanilla'
|
| 19 |
args.arch = 'dino_small_patch16'
|
| 20 |
-
args.device = 'cuda:7'
|
| 21 |
#args.resume = '/fast_scratch/hushell/fluidstack/FS125_few-shot-transformer/outputs/dinosmall_1e-4/best_converted.pth'
|
| 22 |
args.resume = 'https://huggingface.co/hushell/pmf_dinosmall_lr1e-4/resolve/main/best_converted.pth'
|
| 23 |
args.api_key = 'AIzaSyAFkOGnXhy-2ZB0imDvNNqf2rHb98vR_qY'
|
|
@@ -25,7 +24,7 @@ args.cx = '06d75168141bc47f1'
|
|
| 25 |
|
| 26 |
|
| 27 |
# model
|
| 28 |
-
device = torch.device(
|
| 29 |
model = get_model(args)
|
| 30 |
model.to(device)
|
| 31 |
#checkpoint = torch.load(args.resume, map_location='cpu')
|
|
|
|
| 17 |
args = DotMap()
|
| 18 |
args.deploy = 'vanilla'
|
| 19 |
args.arch = 'dino_small_patch16'
|
|
|
|
| 20 |
#args.resume = '/fast_scratch/hushell/fluidstack/FS125_few-shot-transformer/outputs/dinosmall_1e-4/best_converted.pth'
|
| 21 |
args.resume = 'https://huggingface.co/hushell/pmf_dinosmall_lr1e-4/resolve/main/best_converted.pth'
|
| 22 |
args.api_key = 'AIzaSyAFkOGnXhy-2ZB0imDvNNqf2rHb98vR_qY'
|
|
|
|
| 24 |
|
| 25 |
|
| 26 |
# model
|
| 27 |
+
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
|
| 28 |
model = get_model(args)
|
| 29 |
model.to(device)
|
| 30 |
#checkpoint = torch.load(args.resume, map_location='cpu')
|