Python Backend: "Expected all tensors to be on the same device, but found at least two devices, cpu and cuda:0"

Issue Description

I found this bug with the python backend while trying to get it working with ROCm for my AMD GPU, though I imagine the same issue may be present for CUDA too. If PyTorch sees a GPU, then pb2torch in model.py creates tensors on the CPU without transferring it over to the GPU. torch2pb also doesn't transfer tensors back from the GPU before converting for numpy. Failing to do so raises the error in the title.

A workaround in my case is to always transfer to cuda in model.py, though I imagine that would break CPU compatibility

def pb2torch(request, name):
    tensor = pb_utils.get_input_tensor_by_name(request, name)
-    return torch.from_numpy(tensor.as_numpy())
+    return torch.from_numpy(tensor.as_numpy()).to("cuda")
    # return from_dlpack(tensor.to_dlpack())


def torch2pb(name, tensor):
-    return pb_utils.Tensor(name, tensor.numpy())
+    return pb_utils.Tensor(name, tensor.cpu().numpy())
    # return pb_utils.Tensor.from_dlpack(name, to_dlpack(tensor))

How to Reproduce

1. Adjust triton to use ROCm's version of PyTorch (as per discussion #81)
2. Try to run any model

Further Information

• PyTorch 2.1.0 for ROCm 5.6