Fix missing import

Signed-off-by: Nicolas Castet <ncastet@nvidia.com>
horovod · Apr 17, 2023 · ff7052c · ff7052c
1 parent eb8e477
commit ff7052c
Showing 1 changed file with 2 additions and 0 deletions.
diff --git a/test/parallel/test_tensorflow_keras.py b/test/parallel/test_tensorflow_keras.py
@@ -29,11 +29,13 @@
 from horovod.common.util  import is_version_greater_equal_than
 
 if is_version_greater_equal_than(tf.__version__, "2.6.0"):
+    from keras import backend as K
     if version.parse(tf.__version__.replace("-tf", "+tf")) < version.parse("2.9.0"):
         from keras.optimizer_v2.optimizer_v2 import OptimizerV2 as Optimizer
     else:
         from tensorflow.keras.optimizers import Optimizer
 else:
+    from tensorflow.python.keras import backend as K
     from tensorflow.python.keras.optimizer_v2.optimizer_v2 import OptimizerV2 as Optimizer
 
 import horovod.tensorflow.keras as hvd