Update Experiments/Resnet50_classification.py

Experiments/Resnet50_classification.py

@@ -91,7 +91,7 @@ def predict(features_path,image):
     image_tensor = transform(pil_image).unsqueeze(0) 
     resnet = models.resnet50(pretrained=True)
     model_check = HiddenLayer(resnet)
-    model_check.load_state_dict(torch.load("CIFAR_end_hll.pt"))
+    model_check.load_state_dict(torch.load("CIFAR_end_hll.pt",map_location=torch.device('cuda' if torch.cuda.is_available() else 'cpu')))
     model_check.eval()
     
     with torch.no_grad():