hushell commited on
Commit
f138c91
1 Parent(s): aeacea0
Files changed (1) hide show
  1. app.py +1 -2
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(args.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')