Thanks for letting me know. Yeah, what I’ve got for now is a with_grad param in the forward method of my q_model that does what I showed. I wasn’t sure if there was a way to prevent the grad propagating through the graph without setting state like that, but it seems like there isn’t so this will have to do. I would still prefer a declarative rather than procedural API here, but it’s obviously not important.