From ab35653b5ad5c39cc72670814c516d6a8b8431ca Mon Sep 17 00:00:00 2001
From: Duo <50307526+iProzd@users.noreply.github.com>
Date: Thu, 29 Feb 2024 12:32:33 +0800
Subject: [PATCH] Add fparam and aparam

---
 deepmd/pt/train/training.py  | 2 ++
 deepmd/utils/env_mat_stat.py | 5 +++++
 2 files changed, 7 insertions(+)

diff --git a/deepmd/pt/train/training.py b/deepmd/pt/train/training.py
index 55609d6dc3..6b62282817 100644
--- a/deepmd/pt/train/training.py
+++ b/deepmd/pt/train/training.py
@@ -861,6 +861,8 @@ def get_data(self, is_train=True, task_key="Default"):
             "atype",
             "box",
             "spin",
+            "fparam",
+            "aparam",
         ]
         input_dict = {item_key: None for item_key in input_keys}
         label_dict = {}
diff --git a/deepmd/utils/env_mat_stat.py b/deepmd/utils/env_mat_stat.py
index 2fa497b9b6..217c46844b 100644
--- a/deepmd/utils/env_mat_stat.py
+++ b/deepmd/utils/env_mat_stat.py
@@ -1,4 +1,5 @@
 # SPDX-License-Identifier: LGPL-3.0-or-later
+import logging
 from abc import (
     ABC,
     abstractmethod,
@@ -19,6 +20,8 @@
     DPPath,
 )
 
+log = logging.getLogger(__name__)
+
 
 class StatItem:
     """A class to store the statistics of the environment matrix.
@@ -170,10 +173,12 @@ def load_or_compute_stats(
         """
         if path is not None and path.is_dir():
             self.load_stats(path)
+            log.info(f"Load stats from {path}.")
         else:
             self.compute_stats(data)
             if path is not None:
                 self.save_stats(path)
+                log.info(f"Save stats to {path}.")
 
     def get_avg(self, default: float = 0) -> Dict[str, float]:
         """Get the average of the environment matrix.