Improve doc

Author: stefanradev93

bayesflow/datasets/disk_dataset.py

@@ -60,6 +60,11 @@ def __init__(
         augmentations : dict of str to Callable or Callable, optional
             Dictionary of augmentation functions to apply to each corresponding key in the batch
             or a function to apply to the entire batch (possibly adding new keys).
+
+            If you provide a dictionary of functions, each function should accept one element
+            of your output batch and return the corresponding transformed element. Otherwise,
+            your function should accept the entire dictionary output and return a dictionary.
+
             Note - augmentations are applied before the adapter is called and are generally
             transforms that you only want to apply during training.
         **kwargs

bayesflow/datasets/offline_dataset.py

@@ -44,6 +44,11 @@ def __init__(
         augmentations : dict of str to Callable or Callable, optional
             Dictionary of augmentation functions to apply to each corresponding key in the batch
             or a function to apply to the entire batch (possibly adding new keys).
+
+            If you provide a dictionary of functions, each function should accept one element
+            of your output batch and return the corresponding transformed element. Otherwise,
+            your function should accept the entire dictionary output and return a dictionary.
+
             Note - augmentations are applied before the adapter is called and are generally
             transforms that you only want to apply during training.
         **kwargs

bayesflow/datasets/online_dataset.py

@@ -41,6 +41,11 @@ def __init__(
         augmentations : dict of str to Callable or Callable, optional
             Dictionary of augmentation functions to apply to each corresponding key in the batch
             or a function to apply to the entire batch (possibly adding new keys).
+
+            If you provide a dictionary of functions, each function should accept one element
+            of your output batch and return the corresponding transformed element. Otherwise,
+            your function should accept the entire dictionary output and return a dictionary.
+
             Note - augmentations are applied before the adapter is called and are generally
             transforms that you only want to apply during training.
         **kwargs