Hi Thomas,
So I had a chance to implement this and test it against the original version. The two versions seem to perform identically, both in terms of outputs and runtime. Both are very slow with the forward pass and backward pass accounting for the majority of the runtime.
So my questions are:

Is it possible to avoid the compute_f_loss().backward() call by somehow reusing f_grad? (This would avoid a second backward pass.)
Is it possible to reuse the first call to compute_f_loss() so that the forward pass only needs to be done once?

In principle this should be possible.