From 2e79dc40515ea50f904994153391f67ed0130758 Mon Sep 17 00:00:00 2001 From: Han Wang Date: Tue, 26 Sep 2023 13:28:02 +0800 Subject: [PATCH] fix the issue of applying modifier multiple times when batch set is load only once --- deepmd/utils/data.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/deepmd/utils/data.py b/deepmd/utils/data.py index 57bea00fac..423745cddf 100644 --- a/deepmd/utils/data.py +++ b/deepmd/utils/data.py @@ -265,8 +265,6 @@ def get_batch(self, batch_size: int) -> dict: self._load_batch_set(self.train_dirs[self.set_count % self.get_numb_set()]) self.set_count += 1 set_size = self.batch_set["coord"].shape[0] - if self.modifier is not None: - self.modifier.modify_data(self.batch_set, self) iterator_1 = self.iterator + batch_size if iterator_1 >= set_size: iterator_1 = set_size @@ -410,6 +408,8 @@ def _get_subdata(self, data, idx=None): def _load_batch_set(self, set_name: DPPath): if not hasattr(self, "batch_set") or self.get_numb_set() > 1: self.batch_set = self._load_set(set_name) + if self.modifier is not None: + self.modifier.modify_data(self.batch_set, self) self.batch_set, _ = self._shuffle_data(self.batch_set) self.reset_get_batch()