Updated to load model on GPU

app.py

@@ -182,7 +182,7 @@ def load_model():
     # Load the pre-trained CLIP model
     clip_model, clip_preprocess = clip.load("ViT-B/32", device=device)
     inst = get_instrumented_model(config.model, config.output_class,
-                              config.layer, torch.device('cpu'), use_w=config.use_w)
+                              config.layer, device, use_w=config.use_w)
     return clip_model, inst
 
 # Then, to load your models, call this function: