update doc

pina/data/data_module.py

@@ -16,29 +16,23 @@ from ..collector import Collector
 
 
 class DummyDataloader:
-    """
-    Dataloader used when batch size is ``None``. It returns the entire dataset
-    in a single batch.
-    """
 
     def __init__(self, dataset):
         """
-        Preprare a dataloader object which will return the entire dataset
-        in a single batch. Depending on the number of GPUs, the dataset is
-        managed as follows:
+        Prepare a dataloader object that returns the entire dataset in a single
+        batch. Depending on the number of GPUs, the dataset is managed 
+        as follows:
 
-        - **Distributed Environment** (multiple GPUs):
-            - Divides the dataset across processes using the rank and world
-                size.
-            - Fetches only the portion of data corresponding to the current
-                process.
-        - **Non-Distributed Environment** (single GPU):
-            - Fetches the entire dataset.
+        - **Distributed Environment** (multiple GPUs): Divides dataset across 
+            processes using the rank and world size. Fetches only portion of 
+            data corresponding to the current process.
+        - **Non-Distributed Environment** (single GPU): Fetches the entire 
+            dataset.
 
-        :param dataset: The dataset object to be processed.
-        :type dataset: PinaDataset
+        :param PinaDataset dataset: The dataset object to be processed.
 
-        .. note:: This data loader is used when the batch size is ``None``.
+        .. note:: 
+           This dataloader is used when the batch size is ``None``.
         """
 
         if (
@@ -84,8 +78,10 @@ class Collator:
         Initialize the object, setting the collate function based on whether
         automatic batching is enabled or not.
 
-        :param dict max_conditions_lengths: dict containing the maximum number
-            of data points to consider in a single batch for each condition.
+        :param dict max_conditions_lengths: ``dict`` containing the maximum 
+            number  of data points to consider in a single batch for 
+            each condition.
+        :param bool automatic_batching: Whether to enable automatic batching.
         :param PinaDataset dataset: The dataset where the data is stored.
         """
 

pina/data/dataset.py

@@ -276,7 +276,8 @@ class PinaGraphDataset(PinaDataset):
         :param data: List of items to collate in a single batch.
         :type data: list[Data] | list[Graph]
         :return: Batch object.
-        :rtype: Batch | LabelBatch
+        :rtype: :class:`~torch_geometric.data.Batch` 
+            | :class:`~pina.graph.LabelBatch`
         """
 
         if isinstance(data[0], Data):