diff --git a/configs/_base_/datasets/cifar10_nopad.py b/configs/_base_/datasets/cifar10_nopad.py
index 141e404ea5..16dfec7d0c 100644
--- a/configs/_base_/datasets/cifar10_nopad.py
+++ b/configs/_base_/datasets/cifar10_nopad.py
@@ -1,7 +1,5 @@
-# custom_imports = dict(
-#     imports=['mmcls.datasets.transforms'], allow_failed_imports=False)
 cifar_pipeline = [
-    dict(type='Flip', flip_ratio=0.5, direction='horizontal'),
+    dict(type='Flip', keys=['img'], flip_ratio=0.5, direction='horizontal'),
     dict(type='PackEditInputs')
 ]
 cifar_dataset = dict(
diff --git a/configs/_base_/datasets/imagenet_128.py b/configs/_base_/datasets/imagenet_128.py
index 39310f4bae..88a15db1f0 100644
--- a/configs/_base_/datasets/imagenet_128.py
+++ b/configs/_base_/datasets/imagenet_128.py
@@ -7,7 +7,7 @@
     dict(type='LoadImageFromFile', key='img'),
     dict(type='RandomCropLongEdge', keys=['img']),
     dict(type='Resize', scale=(128, 128), keys=['img'], backend='pillow'),
-    dict(type='Flip', flip_ratio=0.5, direction='horizontal'),
+    dict(type='Flip', keys=['img'], flip_ratio=0.5, direction='horizontal'),
     dict(type='PackEditInputs')
 ]