update trainer

Author: dario-coscia

pina/trainer.py

@@ -22,6 +22,7 @@ def __init__(
         automatic_batching=None,
         num_workers=None,
         pin_memory=None,
+        shuffle=None,
         **kwargs,
     ):
         """
@@ -53,6 +54,8 @@ def __init__(
         :type num_workers: int
         :param pin_memory: Whether to use pinned memory for faster data transfer to GPU. (Default False)
         :type pin_memory: bool
+        :param shuffle: Whether to shuffle the data for training. (Default False)
+        :type pin_memory: bool
 
         :Keyword Arguments:
             The additional keyword arguments specify the training setup
@@ -77,6 +80,10 @@ def __init__(
             check_consistency(pin_memory, int)
         else:
             num_workers = 0
+        if shuffle is not None:
+            check_consistency(shuffle, bool)
+        else:
+            shuffle = False
         if train_size + test_size + val_size + predict_size > 1:
             raise ValueError(
                 "train_size, test_size, val_size and predict_size "
@@ -131,6 +138,7 @@ def __init__(
             automatic_batching,
             pin_memory,
             num_workers,
+            shuffle
         )
 
         # logging
@@ -166,6 +174,7 @@ def _create_datamodule(
         automatic_batching,
         pin_memory,
         num_workers,
+        shuffle
     ):
         """
         This method is used here because is resampling is needed
@@ -196,6 +205,7 @@ def _create_datamodule(
             automatic_batching=automatic_batching,
             num_workers=num_workers,
             pin_memory=pin_memory,
+            shuffle=shuffle
         )
 
     def train(self, **kwargs):