So I have a better understanding why this is failing. I think that in order to avoid side-effects when calling the autograd.functional.jacobian function (which is used within hessian), the devs decide to perform a no-op on each of the input tensors and send resulting tensor to the function. The problem with this is that all nn.Parameters are necessarily leaf nodes, so any attempt to assign or copy data into the network parameters within the function necessarily breaks the edge in the computation graph pointing back to the inputs. Furthermore, since the inputs sent in aren’t the same objects that were passed to .jacobian then you end up never including the inputs to compute_z in the computation graph.
The only semi-solution I’ve found is to simply remove/comment out this line 25 from the pytorch implementation.  This just removes the no-op but unfortunately that’s obviously not viable  My guess is this may also have other consequences, for example if there are multiple branching and merging paths in the computation graph then we may get the correct results?