diff --git a/dan/transforms.py b/dan/transforms.py
index 9c2ed868eebd3cbe8e50879e8846dff703398cb2..0ad4420d58335e3235ce01f1b783d84e52f0bf7f 100644
--- a/dan/transforms.py
+++ b/dan/transforms.py
@@ -6,7 +6,7 @@ from enum import Enum
 from random import randint
 
 import numpy as np
-from albumentations import OneOf, SomeOf
+from albumentations import SomeOf
 from albumentations.augmentations import (
     Affine,
     CoarseDropout,
@@ -16,7 +16,6 @@ from albumentations.augmentations import (
     GaussianBlur,
     GaussNoise,
     Perspective,
-    PiecewiseAffine,
     Sharpen,
     ToGray,
 )
@@ -180,14 +179,7 @@ def get_augmentation_transforms() -> SomeOf:
             GaussianBlur(sigma_limit=2.5),
             GaussNoise(var_limit=50**2),
             ColorJitter(contrast=0.2, brightness=0.2, saturation=0.2, hue=0.2),
-            OneOf(
-                [
-                    ElasticTransform(
-                        alpha=20.0, sigma=5.0, alpha_affine=1.0, border_mode=0
-                    ),
-                    PiecewiseAffine(scale=(0.01, 0.04), nb_rows=1, nb_cols=4),
-                ]
-            ),
+            ElasticTransform(alpha=20.0, sigma=5.0, alpha_affine=1.0, border_mode=0),
             Sharpen(alpha=(0.0, 1.0)),
             ErosionDilation(min_kernel=1, max_kernel=4, iterations=1),
             Affine(shear={"x": (-20, 20), "y": (0, 0)}),