Thanks for your reply. x and costate are vectors of shape [7], tof and thrust_mag are scalars of shape [1] and thrust_hat and gravity are vectors of shape [3].
Since my initial post I’ve been able to run the forward() method by itself in a torch.no_grad() context in a Jupyter notebook, so it seems like there’s something special about the odeint_adjoint function that is causing the problem. I might have to dive into the torchdiffeq source code to diagnose the problem, but it’s unclear to me how I could be getting different results calling forward() in a no grad context directly vs. calling odeint_adjoint which then calls forward().