TF 2.4 raises a val err: Keras symbolic inputs/outputs do not impleme…

conda_gpu_py36.yml

@@ -0,0 +1,16 @@
+name: surgeon-py36-tf23-gpu
+
+dependencies:
+  - python<3.7
+  - pip
+  - matplotlib
+  - opencv
+  - cudnn=7
+  - cudatoolkit
+  - pip:
+    - tensorflow-gpu==2.3
+    - numpy<=1.19.5
+    - opencv-python>=4.0
+    - keras
+    - lxml
+    - tqdm

src/kerassurgeon/surgeon.py

@@ -441,15 +441,20 @@ def _apply_delete_mask(self, node, inbound_masks):
                              'MaxPooling3D',
                              'AveragePooling1D', 'AveragePooling2D',
                              'AveragePooling3D'):
-            index = [slice(None, x, None) for x in output_shape[1:]]
-            if data_format == 'channels_first':
-                index[0] = slice(None)
-            elif data_format == 'channels_last':
-                index[-1] = slice(None)
+
+            if output_shape is None:
+                outbound_mask = None
+                new_layer = layer
             else:
-                raise ValueError('Invalid data format')
-            outbound_mask = inbound_masks[tuple(index)]
-            new_layer = layer
+                index = [slice(None, x, None) for x in output_shape[1:]]
+                if data_format == 'channels_first':
+                    index[0] = slice(None)
+                elif data_format == 'channels_last':
+                    index[-1] = slice(None)
+                else:
+                    raise ValueError('Invalid data format')
+                outbound_mask = inbound_masks[tuple(index)]
+                new_layer = layer
 
         elif layer_class in ('UpSampling1D',
                              'UpSampling2D',

src/kerassurgeon/utils.py

@@ -2,6 +2,7 @@
 import numpy as np
 from tensorflow.keras.layers import Layer
 from tensorflow.keras.activations import linear
+from tensorflow.python.keras.engine import keras_tensor
 import tensorflow as tf
 from ._utils import node as node_utils
 
@@ -137,20 +138,30 @@ def sort_x_by_y(x, y):
 
 def single_element(x):
     """If x contains a single element, return it; otherwise return x"""
-    if isinstance(x, tf.Tensor):
+
+    if isinstance(x, (tf.Tensor, keras_tensor.KerasTensor)): 
         return x
 
-    if len(x) == 1:
-        x = x[0]
+    if isinstance(x, list):
+        if len(x) == 1:
+            return x[0]
+
+    if isinstance(x, tuple):
+        return x[0]
+
     return x
 
 
 def get_one_tensor(x):
-    if isinstance(x, tf.Tensor):
+
+    if isinstance(x, (tf.Tensor, keras_tensor.KerasTensor)): 
         return x
 
-    assert len(x) == 1
-    return x[0]
+    if isinstance(x, list):
+        if len(x) == 1:
+            return x[0]
+
+    return x
 
 
 def bool_to_index(x):