fix bugs

Author: FilippoOlivo

pina/_src/data/aggregator.py

@@ -32,6 +32,8 @@ def __len__(self):
         :return: The length of the aggregated dataloader.
         :rtype: int
         """
+        if self.batching_mode == "separate_conditions":
+            return sum(len(dl) for dl in self.dataloaders.values())
         return max(len(dl) for dl in self.dataloaders.values())
 
     def __iter__(self):
@@ -42,10 +44,11 @@ def __iter__(self):
         :rtype: iterator
         """
         if self.batching_mode == "separate_conditions":
-            for name, dl in self.dataloaders.items():
-                for batch in dl:
-                    yield {name: batch}
-            return
+            # TODO: implement separate_conditions batching mode
+            raise NotImplementedError(
+                "Batching mode 'separate_conditions' is not implemented yet."
+            )
+
         iterators = {name: iter(dl) for name, dl in self.dataloaders.items()}
         for _ in range(len(self)):
             batch = {}

pina/_src/data/data_module.py

@@ -7,12 +7,9 @@
 import warnings
 from lightning.pytorch import LightningDataModule
 import torch
-from torch_geometric.data import Data
-from torch.utils.data import DataLoader, SequentialSampler
-from torch.utils.data.distributed import DistributedSampler
-from pina._src.core.label_tensor import LabelTensor
-from pina._src.data.dataset import PinaDatasetFactory, PinaTensorDataset
+from torch_geometric.data import Batch
 from pina._src.data.creator import _Creator
+from pina._src.core.graph import LabelBatch, Graph
 from pina._src.data.aggregator import _Aggregator
 
 
@@ -131,6 +128,7 @@ def __init__(
         self.shuffle = shuffle
         self.batching_mode = batching_mode
         self.automatic_batching = automatic_batching
+        self.batching_mode = batching_mode
 
         # If batch size is None, num_workers has no effect
         if batch_size is None and num_workers != 0:
@@ -244,7 +242,6 @@ def setup(self, stage=None):
         :raises ValueError: If the stage is neither "fit" nor "test".
         """
         if stage == "fit" or stage is None:
-            print("Sono qui")
             self.train_datasets = {
                 name: _ConditionSubset(
                     condition,
@@ -254,7 +251,7 @@ def setup(self, stage=None):
                 for name, condition in self.problem.conditions.items()
                 if len(self.split_idxs[name]["train"]) > 0
             }
-            print(self.train_datasets)
+
             self.val_datasets = {
                 name: _ConditionSubset(
                     condition,
@@ -265,6 +262,7 @@ def setup(self, stage=None):
                 if len(self.split_idxs[name]["val"]) > 0
             }
             return
+
         if stage == "test" or stage is None:
             self.test_datasets = {
                 name: _ConditionSubset(
@@ -281,22 +279,20 @@ def setup(self, stage=None):
             )
 
     def train_dataloader(self):
-        print(self.train_datasets)
         return _Aggregator(
             self.creator(self.train_datasets),
-            batching_mode="separate_conditions",
+            batching_mode=self.batching_mode,
         )
 
     def val_dataloader(self):
-        print(self.val_datasets)
         return _Aggregator(
-            self.creator(self.val_datasets), batching_mode="separate_conditions"
+            self.creator(self.val_datasets), batching_mode=self.batching_mode
         )
 
     def test_dataloader(self):
         return _Aggregator(
             self.creator(self.test_datasets),
-            batching_mode="separate_conditions",
+            batching_mode=self.batching_mode,
         )
 
     @staticmethod