directly load model weight to gpu

utils/load_model.py

@@ -34,7 +34,7 @@ def load_xclip(device: str = "cuda:0",
                     "loss_sym_box_label": 0, "loss_xclip": 0}
     model = OwlViTForClassification(owlvit_det_model=owlvit_det_model, num_classes=n_classes, device=device, weight_dict=weight_dict, logits_from_teacher=use_teacher_logits, custom_box_head=custom_box_head)
     if model_path is not None:
-        ckpt = torch.load(model_path, map_location='cpu')
+        ckpt = torch.load(model_path, map_location=device, weights_only=True)
         model.load_state_dict(ckpt, strict=False)
     model.to(device)
     return model, owlvit_det_processor