How to specify CUDA device number in C++?

Hello,
I have four different libtorch services and also four GPU’s available. In PyTorch its is quite easy to state the GPU device, how do I do it in C++ so that each service is assigned to a different GPU?

Thank,

The .to() operation should also work in C++ as described here.

Thanks, for reference:
Change:

torch::jit::script::Module module = torch::jit::load(s_model_name, torch::kCUDA);

To:

torch::jit::script::Module module = torch::jit::load(s_model_name, torch::kCUDA);
const std::string device_string = "cuda:2";
module.to(device_string);

And:

input_tensor = input_tensor.to(at::kCUDA);

To:

input_tensor = input_tensor.to(device_string);

2 Likes

nice! very simple, thank you very much!