cpu

model_tools.py

@@ -40,7 +40,7 @@ def get_prediction(model, image_in, pipeline):
 
     # Predict
     with torch.no_grad():
-        torch_input = torch.Tensor(preprocessed_image).unsqueeze(0).to('cuda')
+        torch_input = torch.Tensor(preprocessed_image).unsqueeze(0).to('cpu')
         model_output = model(torch_input)
         prediction = pipeline._decode_model_output(model_output, model_input=torch_input)
   # Postprocess