Improvement in DDP and bug fix in DataModule (#432)

2025-01-28 13:51:57 +01:00
parent 629a6ee43b
commit 0194fab0d1
3 changed files with 98 additions and 60 deletions
--- a/pina/trainer.py
+++ b/pina/trainer.py
@@ -72,12 +72,14 @@ class Trainer(lightning.pytorch.Trainer):
            raise RuntimeError('Cannot create Trainer if not all conditions '
                               'are sampled. The Trainer got the following:\n'
                               f'{error_message}')
+        automatic_batching = False
        self.data_module = PinaDataModule(collector=self.solver.problem.collector,
                                          train_size=self.train_size,
                                          test_size=self.test_size,
                                          val_size=self.val_size,
                                          predict_size=self.predict_size,
-                                          batch_size=self.batch_size,)
+                                          batch_size=self.batch_size,
+                                          automatic_batching=automatic_batching)

    def train(self, **kwargs):
        """