Can I use the derivative of part network to the input in the .forward method

Dear all

I am trying to use NN to solve an ODE. The idea is that the trial solution can be written as an expression of NN output. And the loss function can be defined as the square of ODE when substituting the trial solution. Here I met the question that the NN should be derivated by the input during the substitution. And I am wondering how to realize this goal in the definition of my model.

Thank u in advance