fix batch size

Author: FilippoOlivo

pina/_src/data/creator.py

@@ -79,13 +79,14 @@ def _compute_batch_sizes(self, datasets):
         """
         batch_sizes = {}
         if self.batching_mode == "common_batch_size":
+            
+            if self.batch_size is None:
+                batch_size = max(dataset.length for dataset in datasets.values())
+            else:
+                batch_size = self.batch_size
+            
             for name in datasets.keys():
-                if self.batch_size is None:
-                    batch_sizes[name] = len(datasets[name])
-                else:
-                    batch_sizes[name] = min(
-                        self.batch_size, len(datasets[name])
-                    )
+                batch_sizes[name] = min(batch_size, len(datasets[name]))
             return batch_sizes
         if self.batching_mode == "proportional":
             return self._compute_proportional_batch_sizes(datasets)
@@ -168,8 +169,9 @@ def __call__(self, 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":
+            if self.batching_mode == "common_batch_size" and dataset.length != batch_sizes[name]:
                 dataset.max_len = max_len
             dataloaders[name] = self.conditions[name].create_dataloader(
                 dataset=dataset,