title: model.predict doesn't work with Keras Custom Layer (inference error)
body: <p>I've developed a custom convolutional layer. I can use it inside a model and train it (model.fit works), but model.predict() yields an error!</p>
  <p>I will add a simple code to demonstrate how the code is structured.</p>
  <pre><code>modelx1 = tf.keras.models.Sequential([tf.keras.Input(shape=(49,)), Dense(1, activation = 'relu')])
  
  class customLayer(tf.keras.layers.Layer):
   def __init__(self,n=10):super(customLayer, self).__init__()
   def call(self, inputs):
    _, Dim0,Dim1, Dim3 = inputs.shape
    input_victorized = tf.image.extract_patches(images=inputs, sizes=[-1, 7, 7, 1],
                            strides=[1, 1, 1, 1],rates=[1, 1, 1, 1], padding='SAME')
    input_victorized2 = tf.reshape(input_victorized, [-1,49])
    model_output = modelx1(input_victorized2)
    out = tf.reshape(model_output,[-1,Dim0,Dim1,Dim3])
    return out
  </code></pre>
  <p>The custom layer reshapes the input, then feeds it to 'modelx1' then it reshapes the output.</p>
  <p>Here is a simple model where the custom layer is used:</p>
  <pre><code>input1 = tf.keras.Input(shape=(28,28,1))
  x =  Conv2D(filters = 2, kernel_size = 5, activation = 'relu')(input1)
  Layeri = customLayer()(x)
  xxc = Flatten()(Layeri)
  y = Dense(units = 3, activation = 'softmax')(xxc)
  model = tf.keras.Model(inputs=input1, outputs=y)
  model.summary()
  </code></pre>
  <p>The error appears when I run model.predict:</p>
  <pre><code>model.predict(np.ones([100,28,28,1]))
  
  UnimplementedError:  Only support ksizes across space.
       [[node model_58/custom_layer_9/ExtractImagePatches
   (defined at &lt;ipython-input-279-953feb59f882&gt;:7)
  ]] [Op:__inference_predict_function_14640]
  
  Errors may have originated from an input operation.
  Input Source operations connected to node model_58/custom_layer_9/ExtractImagePatches:
  In[0] model_58/conv2d_98/Relu (defined at /usr/local/lib/python3.7/dist-packages/keras/backend.py:4867)
  </code></pre>