Bugfix

keisen · Mar 25, 2024 · 35da52e · 35da52e
1 parent 135aef2
commit 35da52e
Showing 1 changed file with 5 additions and 1 deletion.
diff --git a/tf_keras_vis/utils/__init__.py b/tf_keras_vis/utils/__init__.py
@@ -135,4 +135,8 @@ def lower_precision_dtype(model):
 
 
 def get_input_names(model):
-    return [input.name for input in model.inputs]
+    if version(tf.version.VERSION) >= version("2.4.0"):
+        names = [input.name for input in model.inputs]
+    else:
+        names = model.input_names
+    return names