Repository: mlflow/mlflow

Issue:
[BUG] Predict with Tensor (Numpy) inputs for PyTorch flavor doesn't work on GPUs
### Issues Policy acknowledgement

- [X] I have read and agree to submit bug reports in accordance with the [issues policy](https://www.github.com/mlflow/mlflow/blob/master/ISSUE_POLICY.md)

### Willingness to contribute

Yes. I can contribute a fix for this bug independently.

### MLflow version

- Client: 2.1.1
- Tracking server: 2.1.1


### System information

- **OS Platform and Distribution (e.g., Linux Ubuntu 16.04)**: Linux 20.04
- **Python version**: 3.9
- **yarn version, if running the dev UI**:


### Describe the problem

When running predict on `mlflow.pytorch` models when trying to run inference on a GPU by passing in `device` into the predict call and using numpy array input tensors, prediction fails with the error:

```
TypeError: can't convert cuda:0 device type tensor to numpy. Use Tensor.cpu() to copy the tensor to host memory first.
```

This is because `.numpy` is called on the `preds` tensor [here](https://github.com/mlflow/mlflow/blob/master/mlflow/pytorch/__init__.py#L801), but the `preds` tensor is never brought back to the host CPU (and the scoring server/Python REPL is not running on the GPU). Either you need to call `.detach()` on the `preds` tensor before calling `.numpy` or explicitly pass the tensor back to the CPU.


### Tracking information

<!-- PLEASE KEEP BACKTICKS AND CHECK PREVIEW -->
```shell
REPLACE_ME
```


### Code to reproduce issue

<!-- PLEASE KEEP BACKTICKS AND CHECK PREVIEW -->
```
from PIL import Image
import torch.hub
model = torch.hub.load('pytorch/vision', 'resnet18', pretrained=True)
from torchvision import transforms
import time
filename = 'dog.jpg'
input_image = Image.open(filename)
preprocess = transforms.Compose([
    transforms.Resize(256),
    transforms.CenterCrop(224),
    transforms.ToTensor(),
    transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
])
input_tensor = preprocess(input_image)
input_batch = input_tensor.unsqueeze(0) # create a mini-batch as expected by the model

import os
import requests
import numpy as np
import pandas as pd
import torch
inp_data = input_batch.numpy()
device = 'cuda:0'
model.to(device).eval()
with torch.no_grad():
    input_tensor = torch.from_numpy(inp_data).to(device)
    preds = model(input_tensor).to('cpu')
    predictions = preds.numpy()
print(predictions.shape)
```


### Stack trace

<!-- PLEASE KEEP BACKTICKS AND CHECK PREVIEW -->
```
REPLACE_ME
```


### Other info / logs

<!-- PLEASE KEEP BACKTICKS AND CHECK PREVIEW -->
```
TypeError: can't convert cuda:0 device type tensor to numpy. Use Tensor.cpu() to copy the tensor to host memory first.
```


### What component(s) does this bug affect?

- [ ] `area/artifacts`: Artifact stores and artifact logging
- [ ] `area/build`: Build and test infrastructure for MLflow
- [ ] `area/docs`: MLflow documentation pages
- [ ] `area/examples`: Example code
- [ ] `area/model-registry`: Model Registry service, APIs, and the fluent client calls for Model Registry
- [X] `area/models`: MLmodel format, model serialization/deserialization, flavors
- [ ] `area/recipes`: Recipes, Recipe APIs, Recipe configs, Recipe Templates
- [ ] `area/projects`: MLproject format, project running backends
- [X] `area/scoring`: MLflow Model server, model deployment tools, Spark UDFs
- [ ] `area/server-infra`: MLflow Tracking server backend
- [ ] `area/tracking`: Tracking Service, tracking client APIs, autologging

### What interface(s) does this bug affect?

- [ ] `area/uiux`: Front-end, user experience, plotting, JavaScript, JavaScript dev server
- [X] `area/docker`: Docker use across MLflow's components, such as MLflow Projects and MLflow Models
- [ ] `area/sqlalchemy`: Use of SQLAlchemy in the Tracking Service or Model Registry
- [ ] `area/windows`: Windows support

### What language(s) does this bug affect?

- [ ] `language/r`: R APIs and clients
- [ ] `language/java`: Java APIs and clients
- [ ] `language/new`: Proposals for new client languages

### What integration(s) does this bug affect?

- [X] `integrations/azure`: Azure and Azure ML integrations
- [X] `integrations/sagemaker`: SageMaker integrations
- [X] `integrations/databricks`: Databricks integrations


Discussion:

Resolving patch:
diff --git a/mlflow/pytorch/__init__.py b/mlflow/pytorch/__init__.py
--- a/mlflow/pytorch/__init__.py
+++ b/mlflow/pytorch/__init__.py
@@ -789,6 +789,8 @@ def predict(self, data, device="cpu"):
         with torch.no_grad():
             input_tensor = torch.from_numpy(inp_data).to(device)
             preds = self.pytorch_model(input_tensor)
+            if device != "cpu":
+                preds = preds.to("cpu")
             if not isinstance(preds, torch.Tensor):
                 raise TypeError(
                     "Expected PyTorch model to output a single output tensor, "


Test patch: