Skip to content

Commit fd29c6e

Browse files
authored
Fix updated_options (#111)
* Fix updated_options * Add options init
1 parent f405488 commit fd29c6e

File tree

1 file changed

+4
-0
lines changed

1 file changed

+4
-0
lines changed

src/libtorch.cc

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1885,6 +1885,10 @@ ModelInstanceState::SetInputTensors(
18851885

18861886
const auto torch_dtype =
18871887
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);
18881892

18891893
torch::Tensor input_tensor = torch::from_blob(
18901894
const_cast<char*>(dst_buffer), shape,

0 commit comments

Comments
 (0)