assert systems and batch sizes have the same length

deepmodeling · Mar 28, 2024 · 503c398 · 503c398
1 parent a16d133
commit 503c398
Showing 1 changed file with 1 addition and 0 deletions.
diff --git a/deepmd/pt/utils/dataloader.py b/deepmd/pt/utils/dataloader.py
@@ -125,6 +125,7 @@ def construct_dataset(system):
             self.batch_sizes = batch_size
         else:
             self.batch_sizes = batch_size * np.ones(len(systems), dtype=int)
+        assert len(self.systems) == len(self.batch_sizes)
         for system, batch_size in zip(self.systems, self.batch_sizes):
             if dist.is_initialized():
                 system_sampler = DistributedSampler(system)