fix bug:device constructor

deepmodeling · Feb 3, 2024 · 6c4a641 · 6c4a641
1 parent e61a509
commit 6c4a641
Showing 1 changed file with 1 addition and 3 deletions.
diff --git a/source/api_cc/src/DeepPotPT.cc b/source/api_cc/src/DeepPotPT.cc
@@ -193,11 +193,9 @@ void DeepPotPT::compute(ENERGYVTYPE& ener,
                         const std::vector<VALUETYPE>& coord,
                         const std::vector<int>& atype,
                         const std::vector<VALUETYPE>& box) {
-  torch::Device device;
+  torch::Device device(torch::kCUDA, gpu_id);
   if(cpu_enabled)
     device = torch::Device(torch::kCPU);
-  else
-    device = torch::Device(torch::kCUDA, gpu_id);
   std::vector<VALUETYPE> coord_wrapped = coord;
   int natoms = atype.size();
   auto options = torch::TensorOptions().dtype(torch::kFloat64);