fix bugs

Author: FilippoOlivo

pina/_src/core/trainer.py

@@ -131,12 +131,30 @@ def __init__(
             automatic_batching if automatic_batching is not None else False
         )
 
+        if batch_size is None and batching_mode != "common_batch_size":
+            warnings.warn(
+                "Batching mode is set to "
+                f"{batching_mode} but batch_size is None. "
+                "Batching mode will be set to common_batch_size.",
+                UserWarning,
+            )
+            batching_mode = "common_batch_size"
+
+        if batch_size == 1 and batching_mode == "proportional":
+            warnings.warn(
+                "Batching mode is set to proportional but batch_size is 1. "
+                "Batching mode will be set to common_batch_size.",
+                UserWarning,
+            )
+            batching_mode = "common_batch_size"
+
         # set attributes
         self.compile = compile
         self.solver = solver
         self.batch_size = batch_size
         self._move_to_device()
         self.data_module = None
+
         self._create_datamodule(
             train_size=train_size,
             test_size=test_size,

pina/_src/data/data_module.py

@@ -72,7 +72,7 @@ def __init__(
         val_size=0.1,
         batch_size=None,
         shuffle=True,
-        batching_mode="separate_conditions",
+        batching_mode="common_batch_size",
         automatic_batching=None,
         num_workers=0,
         pin_memory=False,
@@ -95,7 +95,7 @@ def __init__(
             Default ``True``.
         :param str batching_mode: The batching mode to use. Options are
             ``"common_batch_size"``, ``"proportional"``, and
-            ``"separate_conditions"``. Default is ``"separate_conditions"``.
+            ``"separate_conditions"``. Default is ``"common_batch_size"``.
         :param automatic_batching: If ``True``, automatic PyTorch batching
             is performed, which consists of extracting one element at a time
             from the dataset and collating them into a batch. This is useful
@@ -241,7 +241,7 @@ def setup(self, stage=None):
 
         :raises ValueError: If the stage is neither "fit" nor "test".
         """
-        if stage == "fit" or stage is None:
+        if stage in ("fit", None):
             self.train_datasets = {
                 name: _ConditionSubset(
                     condition,
@@ -261,9 +261,8 @@ def setup(self, stage=None):
                 for name, condition in self.problem.conditions.items()
                 if len(self.split_idxs[name]["val"]) > 0
             }
-            return
 
-        if stage == "test" or stage is None:
+        if stage in ("test", None):
             self.test_datasets = {
                 name: _ConditionSubset(
                     condition,
@@ -273,7 +272,7 @@ def setup(self, stage=None):
                 for name, condition in self.problem.conditions.items()
                 if len(self.split_idxs[name]["test"]) > 0
             }
-        else:
+        if stage not in ("fit", "test", None):
             raise ValueError(
                 f"Invalid stage {stage}. Stage must be either 'fit' or 'test'."
             )