Tensorflow gradient of loss with respect to model output gives None

Question

I\'m trying to differentiate my loss function with respect to the model output in the training_step function of a tf.keras.Model. This is my attemp