fix common_batch_size iteration bug

Author: FilippoOlivo

pina/_src/data/creator.py

@@ -166,7 +166,11 @@ def __call__(self, datasets):
         # Compute batch sizes per condition based on batching_mode
         batch_sizes = self._compute_batch_sizes(datasets)
         dataloaders = {}
+        if self.batching_mode == "common_batch_size":
+            max_len = max(len(dataset) for dataset in datasets.values())
         for name, dataset in datasets.items():
+            if self.batching_mode == "common_batch_size":
+                dataset.max_len = max_len
             dataloaders[name] = self.conditions[name].create_dataloader(
                 dataset=dataset,
                 batch_size=batch_sizes[name],

pina/_src/data/data_module.py

@@ -24,9 +24,11 @@ def __init__(self, condition, indices, automatic_batching):
         self.condition = condition
         self.indices = indices
         self.automatic_batching = automatic_batching
+        self.length = len(self.indices)
+        self.max_len = self.length
 
     def __len__(self):
-        return len(self.indices)
+        return self.max_len
 
     def __getitem__(self, idx):
         """
@@ -36,6 +38,8 @@ def __getitem__(self, idx):
         :return: The data corresponding to the given index.
         :rtype: dict
         """
+        if idx >= self.length:
+            idx = idx % self.length
         idx = self.indices[idx]
         if not self.automatic_batching:
             return idx