We read every piece of feedback, and take your input very seriously.
To see all available qualifiers, see our documentation.
There was an error while loading. Please reload this page.
1 parent f405488 commit fd29c6eCopy full SHA for fd29c6e
src/libtorch.cc
@@ -1885,6 +1885,10 @@ ModelInstanceState::SetInputTensors(
1885
1886
const auto torch_dtype =
1887
ConvertDataTypeToTorchType(batch_input.DataType());
1888
+ torch::TensorOptions options{torch_dtype.second};
1889
+ auto updated_options = (dst_memory_type == TRITONSERVER_MEMORY_GPU)
1890
+ ? options.device(torch::kCUDA, device_.index())
1891
+ : options.device(torch::kCPU);
1892
1893
torch::Tensor input_tensor = torch::from_blob(
1894
const_cast<char*>(dst_buffer), shape,
0 commit comments