fix p argument calling

deepmodeling · Jan 28, 2024 · 893d904 · 893d904
1 parent b3f0b8d
commit 893d904
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/deepmd/pt/utils/dataset.py b/deepmd/pt/utils/dataset.py
@@ -879,7 +879,7 @@ def __len__(self):
     def __getitem__(self, index=None):
         """Get a batch of frames from the selected system."""
         if index is None:
-            index = dp_random.choice(np.arange(self.nsystems), self.probs)
+            index = dp_random.choice(np.arange(self.nsystems), p=self.probs)
         b_data = self._data_systems[index].get_batch(self._batch_size)
         b_data["natoms"] = torch.tensor(
             self._natoms_vec[index], device=env.PREPROCESS_DEVICE
@@ -892,7 +892,7 @@ def __getitem__(self, index=None):
     def get_training_batch(self, index=None):
         """Get a batch of frames from the selected system."""
         if index is None:
-            index = dp_random.choice(np.arange(self.nsystems), self.probs)
+            index = dp_random.choice(np.arange(self.nsystems), p=self.probs)
         b_data = self._data_systems[index].get_batch_for_train(self._batch_size)
         b_data["natoms"] = torch.tensor(
             self._natoms_vec[index], device=env.PREPROCESS_DEVICE