diff --git a/mu_map/training/default.py b/mu_map/training/default.py
index fe8aa382f1bb32503f93eee4b947bdb1dd3e3333..fd554ffd49b9c25c88b0a70855a38433b4db3048 100644
--- a/mu_map/training/default.py
+++ b/mu_map/training/default.py
@@ -158,6 +158,29 @@ if __name__ == "__main__":
         default="mean",
         help="type of normalization applied to the reconstructions",
     )
+    parser.add_argument(
+        "--patch_size",
+        type=int,
+        default=32,
+        help="the size of patches extracted for each reconstruction"
+    )
+    parser.add_argument(
+        "--patch_offset",
+        type=int,
+        default=20,
+        help="offset to ignore the border of the image"
+    )
+    parser.add_argument(
+        "--number_of_patches",
+        type=int,
+        default=100,
+        help="number of patches extracted for each image"
+    )
+    parser.add_argument(
+        "--no_shuffle",
+        action="store_true" 
+        help="do not shuffle patches in the dataset"
+    )
 
     # Training Args
     parser.add_argument(
@@ -265,6 +288,10 @@ if __name__ == "__main__":
     for split in ["train", "validation"]:
         dataset = MuMapPatchDataset(
             args.dataset_dir,
+            patches_per_image=args.number_of_patches,
+            patch_size=args.patch_size,
+            patch_offset=args.patch_offset,
+            shuffle=not args.no_shuffle,
             split_name=split,
             transform_normalization=transform_normalization,
             transform_augmentation=transform_augmentation,