fix bug: device constructor

deepmodeling · Feb 3, 2024 · e61a509 · e61a509
1 parent 7bb1b1c
commit e61a509
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
@@ -71,11 +71,9 @@ void DeepPotPT::compute(ENERGYVTYPE& ener,
                         const std::vector<VALUETYPE>& box,
                         const InputNlist& lmp_list,
                         const int& ago) {
-  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);