clean the code

deepmodeling · Feb 13, 2025 · 036a9df · 036a9df
1 parent e943950
commit 036a9df
Showing 1 changed file with 3 additions and 5 deletions.
diff --git a/source/tests/pt/test_make_stat_input.py b/source/tests/pt/test_make_stat_input.py
@@ -121,7 +121,7 @@ def test_with_missing_elements_and_new_frames(self):
         lst = make_stat_input(
             datasets=self.datasets,
             dataloaders=self.dataloaders,
-            nbatches=10,
+            nbatches=1,
             min_frames_per_element_forstat=1,
             enable_element_completion=False,
         )
@@ -131,21 +131,19 @@ def test_with_missing_elements_and_new_frames(self):
             if "energy" in sys_stat:
                 energy = sys_stat["energy"]
                 missing_elements.append(self.count_non_zero_elements(energy))
-
-        #
+
         self.assertGreater(
             len(missing_elements), 0, "Expected missing elements to be processed."
         )
 
         lst_new = make_stat_input(
             datasets=self.datasets,
             dataloaders=self.dataloaders,
-            nbatches=10,
+            nbatches=1,
             min_frames_per_element_forstat=1,
             enable_element_completion=True,
         )
 
-        #
         for original, new in zip(lst, lst_new):
             energy_ori = np.array(original["energy"].cpu()).flatten()
             energy_new = np.array(new["energy"].cpu()).flatten()