It doesn’t seem to work :-(. I get a runtime error where torch complains that he’s expecting a Tensor but finds a None instead in ctx->saved_data in backward during training. When I replace the above pseudo-code with the following the code runs fine:
forward_2(...);
ctx->saved_data["..."] = ...; // save the pre-computed stuff that backward will need

(ie I don’t check whether autograd is on now and I systematically compute what’s needed for backward)
Looks like Gradmode::is_enabled() doesn’t carry the correct value when inside a torch::autograd::Function?