[scripts] Set datamodule batch size

src/mednet/libs/common/scripts/predict.py

@@ -127,6 +127,7 @@ def setup_datamodule(
     parallel,
 ) -> None:  # numpydoc ignore=PR01
     """Configure and set up the datamodule."""
+    datamodule.batch_size = batch_size
     datamodule.parallel = parallel
     datamodule.model_transforms = model.model_transforms
 

src/mednet/libs/common/scripts/train.py

@@ -275,11 +275,13 @@ def load_checkpoint(checkpoint_file, datamodule, model):
 def setup_datamodule(
     datamodule,
     model,
+    batch_size,
     drop_incomplete_batch,
     cache_samples,
     parallel,
 ) -> None:  # numpydoc ignore=PR01
     """Configure and set up the datamodule."""
+    datamodule.batch_size = batch_size
     datamodule.drop_incomplete_batch = drop_incomplete_batch
     datamodule.cache_samples = cache_samples
     datamodule.parallel = parallel

src/mednet/libs/segmentation/scripts/train.py

@@ -61,6 +61,7 @@ def train(
     setup_datamodule(
         datamodule,
         model,
+        batch_size,
         drop_incomplete_batch,
         cache_samples,
         parallel,