fail to train a model using tf.function

Question

I tried to use tf.function to decorate the gradient update function as below.

import tensorflow as tf
from tensorflow.keras import layers, activation