Solved by Brunobio in post #2 
              
              
                I found the problem, I was losing the gradient at the end of the ode function, because I was  creating a new tensor instead of simply concatenate the tensor like so: 
    def ode(self, x0, r):
        qF = r[0, 0]
        qA = r[0, 1]
        qP = r[0, 2]
        mu = r[0, 3]

        FRU = x0[0, 0]…