@tf.function error when changing optimizers

Question

I\'m trying to use tf.function to handle a train step for a simple model. My train_step function looks like this:

def loss_computati