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);


image.png1185×400 36.1 KB