Use original device for created batch inputs

Author: dyastremsky

src/libtorch.cc

@@ -1884,7 +1884,9 @@ ModelInstanceState::SetInputTensors(
       const auto torch_dtype =
           ConvertDataTypeToTorchType(batch_input.DataType());
       torch::TensorOptions options{torch_dtype.second};
-      auto updated_options = options.device(torch::kCPU);
+      auto updated_options = (memory_type == TRITONSERVER_MEMORY_GPU)
+                               ? options.device(torch::kCUDA, device_.index())
+                               : options.device(torch::kCPU);
 
       torch::Tensor input_tensor = torch::from_blob(
           const_cast<char*>(dst_buffer), shape, updated_options);