Problem with custom loss functions to solve differential equations with Tensorflow/Keras

Question

I\'m trying to reproduce the results from 1902.05563. One can find an example in this link which is solving the eq. 4 in the paper. However, they are writing the layers by h