From 240c54195bb4b437ec26a0de263ab2a3f98b3421 Mon Sep 17 00:00:00 2001
From: Chongxiao Cao <chongxiaoc@uber.com>
Date: Thu, 24 Feb 2022 22:12:31 -0800
Subject: [PATCH] Change the way of import optimizer_v2 for Keras>=2.9.0

---
 test/parallel/test_tensorflow2_keras.py | 5 ++++-
 1 file changed, 4 insertions(+), 1 deletion(-)

diff --git a/test/parallel/test_tensorflow2_keras.py b/test/parallel/test_tensorflow2_keras.py
index d15fc3ad8d..7647b6916d 100644
--- a/test/parallel/test_tensorflow2_keras.py
+++ b/test/parallel/test_tensorflow2_keras.py
@@ -30,7 +30,10 @@
 from horovod.common.util  import is_version_greater_equal_than
 
 if is_version_greater_equal_than(tf.__version__, "2.6.0"):
-    from keras.optimizer_v2 import optimizer_v2
+    if LooseVersion(keras.__version__) < LooseVersion("2.9.0"):
+        from keras.optimizer_v2 import optimizer_v2
+    else:
+        from keras.optimizers.optimizer_v2 import optimizer_v2
 else:
     from tensorflow.python.keras.optimizer_v2 import optimizer_v2