tensorflow: Bug in tf,keras.layers.Conv2D when use dilation_rate

Please make sure that this is a bug. As per our GitHub Policy, we only address code/doc bugs, performance issues, feature requests and build/installation issues on GitHub. tag:bug_template

import tensorflow as tf # TF2

class Conv2D_BN_ReLU(tf.keras.Model):
    """Conv2D + BN + ReLU"""
    def __init__(self,
            filters,
            kernel_size,
            strides=1,
            padding="SAME",
            dilation_rate=1,
            use_bias=False,
            kernel_initializer="he_normal",
            kernel_regularizer=None,
            **bn_params):
        super(Conv2D_BN_ReLU, self).__init__()

        self.conv = tf.keras.layers.Conv2D(filters, kernel_size, strides=strides,
                padding=padding, dilation_rate=dilation_rate, use_bias=use_bias,
                kernel_initializer=kernel_initializer, kernel_regularizer=kernel_regularizer)
        self.bn = tf.keras.layers.BatchNormalization(**bn_params)

    def call(self, x, training=None):
        x = self.conv(x)
        x = tf.nn.relu(self.bn(x, training=training))
        return x

if __name__ == "__main__":
    import os
    os.environ["CUDA_VISIBLE_DEVICES"] = "6"

    net = Conv2D_BN_ReLU(64, 1, 1, dilation_rate=6)
    x = tf.ones((1, 32, 32, 64))
    y = net(x)

    x = tf.ones((1, 64, 64, 64))
    y = net(x)

System information

• Have I written custom code (as opposed to using a stock example script provided in TensorFlow):
• OS Platform and Distribution (e.g., Linux Ubuntu 16.04):
• Mobile device (e.g. iPhone 8, Pixel 2, Samsung Galaxy) if the issue happens on mobile device:
• TensorFlow installed from (source or binary):
• TensorFlow version (use command below): 2.0
• Python version: 3.6
• Bazel version (if compiling from source):
• GCC/Compiler version (if compiling from source):
• CUDA/cuDNN version:
• GPU model and memory:

You can collect some of this information using our environment capture script You can also obtain the TensorFlow version with: 1. TF 1.0: python -c "import tensorflow as tf; print(tf.GIT_VERSION, tf.VERSION)" 2. TF 2.0: python -c "import tensorflow as tf; print(tf.version.GIT_VERSION, tf.version.VERSION)"

Describe the current behavior

tensorflow.python.framework.errors_impl.InvalidArgumentError: padded_shape[0]=68 is not divisible by block_shape[0]=6 [Op:SpaceToBatchND]

Describe the expected behavior

I think the second evaluation should also work, because I use padding=“SAME”

Code to reproduce the issue Provide a reproducible test case that is the bare minimum necessary to generate the problem.

Other info / logs Include any logs or source code that would be helpful to diagnose the problem. If including tracebacks, please include the full traceback. Large logs and files should be attached.