fix(pt): set device for PT C++

Fix deepmodeling#4171. Signed-off-by: Jinzhe Zeng <[email protected]>
njzjz · Oct 26, 2024 · 5f7707a · 5f7707a
1 parent fa61d69
commit 5f7707a
Showing 1 changed file with 1 addition and 0 deletions.
diff --git a/source/api_cc/src/DeepPotPT.cc b/source/api_cc/src/DeepPotPT.cc
@@ -80,6 +80,7 @@ void DeepPotPT::init(const std::string& model,
     device = torch::Device(torch::kCPU);
     std::cout << "load model from: " << model << " to cpu " << std::endl;
   } else {
+    c10::cuda::set_device(gpu_id);
     std::cout << "load model from: " << model << " to gpu " << gpu_id
               << std::endl;
   }