make mapping not none during lmp steps

deepmodeling · Oct 14, 2024 · f3d0b1e · f3d0b1e
1 parent 8279cca
commit f3d0b1e
Show file tree

Hide file tree

Showing 2 changed files with 1 addition and 1 deletion.
diff --git a/source/api_cc/include/DeepPotPT.h b/source/api_cc/include/DeepPotPT.h
@@ -338,6 +338,7 @@ class DeepPotPT : public DeepPotBase {
   int do_message_passing;  // 1:dpa2 model 0:others
   bool gpu_enabled;
   at::Tensor firstneigh_tensor;
+  c10::optional<torch::Tensor> mapping_tensor;
   torch::Dict<std::string, torch::Tensor> comm_dict;
   /**
    * @brief Translate PyTorch exceptions to the DeePMD-kit exception.

diff --git a/source/api_cc/src/DeepPotPT.cc b/source/api_cc/src/DeepPotPT.cc
@@ -164,7 +164,6 @@ void DeepPotPT::compute(ENERGYVTYPE& ener,
   std::vector<std::int64_t> atype_64(datype.begin(), datype.end());
   at::Tensor atype_Tensor =
       torch::from_blob(atype_64.data(), {1, nall_real}, int_option).to(device);
-  c10::optional<torch::Tensor> mapping_tensor;
   if (ago == 0) {
     nlist_data.copy_from_nlist(lmp_list);
     nlist_data.shuffle_exclude_empty(fwd_map);