Testing device=cuda

app.py

@@ -43,7 +43,7 @@ def load_model_hf(model_config_path, repo_id, filename, device='cpu'):
     args.device = device
 
     cache_file = hf_hub_download(repo_id=repo_id, filename=filename)
-    checkpoint = torch.load(cache_file, map_location='cpu')
+    checkpoint = torch.load(cache_file, map_location=device)
     log = model.load_state_dict(clean_state_dict(checkpoint['model']), strict=False)
     print("Model loaded from {} \n => {}".format(cache_file, log))
     _ = model.eval()
@@ -65,7 +65,7 @@ def image_transform_grounding_for_vis(init_image):
     image, _ = transform(init_image, None) # 3, h, w
     return image
 
-model = load_model_hf(config_file, ckpt_repo_id, ckpt_filenmae)
+model = load_model_hf(config_file, ckpt_repo_id, ckpt_filenmae, device=device)
 
 def run_grounding(input_image, grounding_caption, box_threshold, text_threshold):
     init_image = input_image.convert("RGB")