TensorFlow RuntimeError: “Attempting to capture an EagerTensor without building a function”

Question

I am trying to build a neural network in Python for solving PDEs, and, as such, I have had to write custom training steps. My training function looks like this: