diff --git a/the_well/data/datasets.py b/the_well/data/datasets.py
index c3b82593..7620c788 100644
--- a/the_well/data/datasets.py
+++ b/the_well/data/datasets.py
@@ -136,6 +136,8 @@ class WellDataset(Dataset):
             Whether to normalize data in the dataset
         normlization_type:
             What type of dataset normalization. Callable Options: ZSCORE and RMS
+        max_rollout_steps:
+            Maximum number of output steps to return in a single sample. Return the full trajectory if larger than its actual length.
         n_steps_input:
             Number of steps to include in each sample
         n_steps_output: