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): colab
• Mobile device (e.g. iPhone 8, Pixel 2, Samsung Galaxy) if the issue happens on mobile device:
• TensorFlow installed from (source or binary): binary
• TensorFlow version (use command below): 2.4.0rc4
• Python version: any
• Bazel version (if compiling from source): N/A
• GCC/Compiler version (if compiling from source): N/A
• CUDA/cuDNN version: N/A
• GPU model and memory: N/A

v2.4.0-rc3-20-g97c3fef64ba 2.4.0-rc4

Describe the current behavior

RMSprop failes if using mixed precision training and eager functions

Describe the expected behavior

succeed without errors

Standalone code to reproduce the issue

https://colab.research.google.com/drive/1-ZrJyKtaWmOMtga44Efc8DOu1fnRo1wN?usp=sharing

Originated from addons tensorflow/addons#2250