Add compatibility with Keras export script (#455)

* Add compatibility with Keras export script
larq · Aug 7, 2023 · 321ba2c · 321ba2c
1 parent 5e8b37a
commit 321ba2c
Showing 1 changed file with 5 additions and 1 deletion.
diff --git a/larq_zoo/core/utils.py b/larq_zoo/core/utils.py
@@ -14,7 +14,11 @@
     from tensorflow.python.keras.backend import is_keras_tensor
 
 if version.parse(tf.__version__) >= version.parse("2.6"):
-    from keras.applications.imagenet_utils import obtain_input_shape
+    try:
+        from keras.applications.imagenet_utils import obtain_input_shape
+    except ImportError:
+        # type: ignore
+        from keras.src.applications.imagenet_utils import obtain_input_shape
 else:
     try:
         # type: ignore