diff --git a/miditok/constants.py b/miditok/constants.py
index 948ad20f..37c2fc77 100644
--- a/miditok/constants.py
+++ b/miditok/constants.py
@@ -2,7 +2,7 @@
 
 """
 
-CURRENT_VERSION_PACKAGE = "2.0.3"  # used when saving the config of a tokenizer
+CURRENT_VERSION_PACKAGE = "2.0.4"  # used when saving the config of a tokenizer
 
 MIDI_FILES_EXTENSIONS = [".mid", ".midi", ".MID", ".MIDI"]
 
diff --git a/miditok/data_augmentation/data_augmentation.py b/miditok/data_augmentation/data_augmentation.py
index 2469878c..930c02bd 100644
--- a/miditok/data_augmentation/data_augmentation.py
+++ b/miditok/data_augmentation/data_augmentation.py
@@ -92,10 +92,13 @@ def data_augmentation_dataset(
             for track, (_, is_drum) in zip(ids, programs):
                 if is_drum:  # we dont augment drums
                     continue
+                corrected_offsets = deepcopy(offsets)
+                vel_dim = int(128 / len(tokenizer.velocities))
+                corrected_offsets[1] = [int(off / vel_dim) for off in corrected_offsets[1]]
                 aug = data_augmentation_tokens(
                     np.array(track),
                     tokenizer,
-                    *offsets,
+                    *corrected_offsets,
                     all_offset_combinations=all_offset_combinations,
                 )
                 if len(aug) == 0:
@@ -220,7 +223,7 @@ def get_offsets(
             as a tuple of two booleans. (default: (True, True))
     :param midi: midi object to augment (default: None)
     :param ids: token ids as a list of tracks (default: None)
-    :return: augmented MIDI objects.
+    :return: the offsets of pitch, velocity and duration features, in "absolute" value
     """
     offsets = []
 
diff --git a/setup.py b/setup.py
index 03566976..deaae353 100644
--- a/setup.py
+++ b/setup.py
@@ -8,7 +8,7 @@
     author='Nathan Fradet',
     url='https://github.com/Natooz/MidiTok',
     packages=find_packages(exclude=("tests",)),
-    version='2.0.3',
+    version='2.0.4',
     license='MIT',
     description='A convenient MIDI tokenizer for Deep Learning networks, with multiple encoding strategies',
     long_description=long_description,