cpu fixed

pipeline.py

@@ -23,11 +23,11 @@ class Pipeline:
     
     def _load_models(self):
         classification_model = resnet_model
-        classification_model.load_state_dict(torch.load('weights/classification_models/resnet50.pt'))
+        classification_model.load_state_dict(torch.load('weights/classification_models/resnet50.pt', map_location=torch.device('cpu')))
         classification_model.eval()
 
         segmentation_model = ResNetUnet()
-        checkpoint = torch.load('weights/segmentation_models/ResNetUnet_best.pt')
+        checkpoint = torch.load('weights/segmentation_models/ResNetUnet_best.pt', map_location=torch.device('cpu'))
         segmentation_model.load_state_dict(checkpoint['model_state_dict'])
         segmentation_model.eval()