fix(pt): optimize createNlistTensor (#4403)

<!-- This is an auto-generated comment: release notes by coderabbit.ai
-->

## Summary by CodeRabbit

- **New Features**
- Enhanced tensor creation process for improved performance and
efficiency.
  
- **Bug Fixes**
- Improved error handling for PyTorch-related exceptions, providing
clearer error messages.

<!-- end of auto-generated comment: release notes by coderabbit.ai -->

Signed-off-by: Jinzhe Zeng <[email protected]>

source/api_cc/src/DeepPotPT.cc

@@ -31,20 +31,20 @@ void DeepPotPT::translate_error(std::function<void()> f) {
 }
 
 torch::Tensor createNlistTensor(const std::vector<std::vector<int>>& data) {
-  std::vector<torch::Tensor> row_tensors;
-
+  size_t total_size = 0;
   for (const auto& row : data) {
-    torch::Tensor row_tensor = torch::tensor(row, torch::kInt32).unsqueeze(0);
-    row_tensors.push_back(row_tensor);
+    total_size += row.size();
   }
-
-  torch::Tensor tensor;
-  if (row_tensors.size() > 0) {
-    tensor = torch::cat(row_tensors, 0).unsqueeze(0);
-  } else {
-    tensor = torch::empty({1, 0, 0}, torch::kInt32);
+  std::vector<int> flat_data;
+  flat_data.reserve(total_size);
+  for (const auto& row : data) {
+    flat_data.insert(flat_data.end(), row.begin(), row.end());
   }
-  return tensor;
+
+  torch::Tensor flat_tensor = torch::tensor(flat_data, torch::kInt32);
+  int nloc = data.size();
+  int nnei = nloc > 0 ? total_size / nloc : 0;
+  return flat_tensor.view({1, nloc, nnei});
 }
 DeepPotPT::DeepPotPT() : inited(false) {}
 DeepPotPT::DeepPotPT(const std::string& model,

source/api_cc/src/DeepSpinPT.cc

@@ -31,20 +31,20 @@ void DeepSpinPT::translate_error(std::function<void()> f) {
 }
 
 torch::Tensor createNlistTensor2(const std::vector<std::vector<int>>& data) {
-  std::vector<torch::Tensor> row_tensors;
-
+  size_t total_size = 0;
   for (const auto& row : data) {
-    torch::Tensor row_tensor = torch::tensor(row, torch::kInt32).unsqueeze(0);
-    row_tensors.push_back(row_tensor);
+    total_size += row.size();
   }
-
-  torch::Tensor tensor;
-  if (row_tensors.size() > 0) {
-    tensor = torch::cat(row_tensors, 0).unsqueeze(0);
-  } else {
-    tensor = torch::empty({1, 0, 0}, torch::kInt32);
+  std::vector<int> flat_data;
+  flat_data.reserve(total_size);
+  for (const auto& row : data) {
+    flat_data.insert(flat_data.end(), row.begin(), row.end());
   }
-  return tensor;
+
+  torch::Tensor flat_tensor = torch::tensor(flat_data, torch::kInt32);
+  int nloc = data.size();
+  int nnei = nloc > 0 ? total_size / nloc : 0;
+  return flat_tensor.view({1, nloc, nnei});
 }
 DeepSpinPT::DeepSpinPT() : inited(false) {}
 DeepSpinPT::DeepSpinPT(const std::string& model,