DropBlock2D not working with TensorFlow backend in graph mode on Keras 3

[Lescurel]

Current Behavior:

On Keras 3 with TensorFlow 2.16.1, DropBlock2D fails on graph mode with the following casting error:

TypeError: Exception encountered when calling DropBlock2D.call().

int() argument must be a string, a bytes-like object or a real number, not 'SymbolicTensor'

Arguments received by DropBlock2D.call():
  • x=tf.Tensor(shape=(32, 8, 8, 8), dtype=float32)
  • training=True

Steps To Reproduce:

See the following collab:
DropBlock2D Graph Mode

Version:

keras_cv.__version__='0.9.0'
keras.__version__='3.0.0'
tf.__version__='2.16.1'

Anything Else:

The fix is trivial, just use ops.cast() instead of python int() in the call method of the Layer.

[mehtamansi29]

Hi @Lescurel -

Thanks for reporting the issue. I have tested the code snippet and reproduces the reported behaviour. Attached gist file for reference.
We will look into the issue and update you the same.