Hi,

1a) The output produced by that layer is then saved for autograd on the later backwards pass.

This is not strictly true I think. In particular it depends what you call a “layer”.
But the layers in torch.nn and neural nets are completely independent of the autograd. So this statement would be wrong in that case.
The autograd works at a lower level where we have a set of specific “elementary ops” for which we do (only when needed, not all the time!) save input and/or outputs.
Are you interested in the first or second in your case?
I can’t really think of an easy way to hack this in though. Maybe you can use this 3 construct that is used to store Tensors to be used in the backward. This is created at the same time (sometimes before, sometimes after) as the forward function and the unpack is called during the backward pass when the Tensor is needed.
Hope this helps!