Update make_predictions.py

make_predictions.py

@@ -11,7 +11,7 @@ def make_predictions(model, input_img, threshold = 0.05):
         image = cv2.resize(image, (256, 256))
         image = np.transpose(image, (2, 0, 1))
         image = np.expand_dims(image, 0)
-        image = torch.from_numpy(image).to(device)
+        image = torch.from_numpy(image).to('cpu')
         predMask = model(image).squeeze()
         predMask = torch.sigmoid(predMask)
         predMask = predMask.cpu().numpy()