diff --git a/tensorflow_addons/image/dense_image_warp_test.py b/tensorflow_addons/image/dense_image_warp_test.py
index 6810d6a936..872dd3f6b5 100644
--- a/tensorflow_addons/image/dense_image_warp_test.py
+++ b/tensorflow_addons/image/dense_image_warp_test.py
@@ -204,7 +204,7 @@ def test_gradients_exist(self):
         interp = dense_image_warp(image, flows)
         loss = tf.math.reduce_mean(tf.math.square(interp - image))
 
-        optimizer = tf.optimizers.Adam(1.0)
+        optimizer = tf.keras.optimizers.Adam(1.0)
         grad = tf.gradients(loss, [flows])
         opt_func = optimizer.apply_gradients(zip(grad, [flows]))
         init_op = tf.compat.v1.global_variables_initializer()
diff --git a/tensorflow_addons/layers/wrappers_test.py b/tensorflow_addons/layers/wrappers_test.py
index 46cb997eb6..9d83bbec50 100644
--- a/tensorflow_addons/layers/wrappers_test.py
+++ b/tensorflow_addons/layers/wrappers_test.py
@@ -66,7 +66,8 @@ def test_weightnorm_conv2d(self):
 
         model.add(tf.keras.layers.Activation('relu'))
         model.compile(
-            optimizer=tf.optimizers.RMSprop(learning_rate=0.001), loss='mse')
+            optimizer=tf.keras.optimizers.RMSprop(learning_rate=0.001),
+            loss='mse')
         model.fit(
             np.random.random((2, 4, 4, 3)),
             np.random.random((2, 4, 4, 5)),
diff --git a/tensorflow_addons/optimizers/weight_decay_optimizers_test.py b/tensorflow_addons/optimizers/weight_decay_optimizers_test.py
index a4c203975c..fc22aa0c56 100644
--- a/tensorflow_addons/optimizers/weight_decay_optimizers_test.py
+++ b/tensorflow_addons/optimizers/weight_decay_optimizers_test.py
@@ -269,7 +269,7 @@ class ExtendWithWeightDecayTest(SGDWTest):
     """Verify that the factory function SGDW is the same as SGDW."""
 
     optimizer = weight_decay_optimizers.extend_with_decoupled_weight_decay(
-        tf.optimizers.SGD)
+        tf.keras.optimizers.SGD)
 
 
 if __name__ == "__main__":