Change the way of import optimizer_v2 for Keras>=2.9.0

horovod · Feb 25, 2022 · 6e4b8f5 · 6e4b8f5
1 parent 02d62dd
commit 6e4b8f5
Show file tree

Hide file tree

Showing 2 changed files with 8 additions and 2 deletions.
diff --git 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
 

diff --git a/test/parallel/test_tensorflow_keras.py b/test/parallel/test_tensorflow_keras.py
@@ -30,7 +30,10 @@
 
 if is_version_greater_equal_than(tf.__version__, "2.6.0"):
     from keras import backend as K
-    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 import backend as K
     from tensorflow.python.keras.optimizer_v2 import optimizer_v2