diff --git a/experimental/marginalization_mixup/augment.py b/experimental/marginalization_mixup/augment.py
index b0d9f7f0..0c7d2110 100644
--- a/experimental/marginalization_mixup/augment.py
+++ b/experimental/marginalization_mixup/augment.py
@@ -17,7 +17,7 @@
 
 import math
 
-from keras.layers.preprocessing import image_preprocessing as image_ops
+from tensorflow.keras.layers.preprocessing import image_preprocessing as image_ops
 import tensorflow as tf