tensorflow: Conv2D breaking when used with keras Model subclassing

System information

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

Describe the current behavior I am trying to use tf.keras.Model subclassing to define my model but it is failing on the first conv layer. Here is the code I am using:

class CustomModel(tf.keras.Model):
  def __init__(self, **kwargs):
    super(CustomModel, self).__init__(**kwargs)
    self.conv1   = Conv2D(32, (3, 3), padding='same')
    self.conv2   = Conv2D(64, (3, 3), padding='same')
    self.pool    = MaxPooling2D(pool_size=(2, 2))
    self.bn      = BatchNormalization()
    self.relu    = Activation("relu")
    self.softmax = Activation("softmax")
    self.drop1   = Dropout(0.25)
    self.drop2   = Dropout(0.5)
    self.dense1  = Dense(512)
    self.dense2  = Dense(10)
    self.flat    = Flatten()
    
  
  
  def call(self, inputs, train):
    z = self.conv1(inputs)
    z = self.bn(z, training=train)
    z = self.relu(z)
    
    z = self.conv1(z)
    z = self.bn(z, training=train)
    z = self.relu(z)
    z = self.pool(z)
    z = self.drop1(z, training=train)
    
    z = self.conv2(z)
    z = self.bn(z, training=train)
    z = self.relu(z)
    
    z = self.conv2(z)
    z = self.bn(z, training=train)
    z = self.relu(z)
    z = self.pool(z)
    z = self.drop1(z, training=train)
    
    z = self.flat(z)
    z = self.dense1(z)
    z = self.relu(z)
    z = self.drop2(z, training=train)
    z = self.dense2(z)
    z = self.softmax(z)
    
    return z

In order to check if the model is working fine, I am passing a random input to the model in this way:

random_input = np.random.rand(32,32, 3).astype(np.float32)
random_input = np.expand_dims(random_input, axis=0)
preds = model(random_input, train=False)

But this throws the following error: InvalidArgumentError: input depth must be evenly divisible by filter depth: 32 vs 3 [Op:Conv2D]

The same model works fine with Sequential/Functional API. So either I am missing out on something or there is something wrong with the subclassing. Can you please look into it?