Update model_tools.py

model_tools.py

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