I think the problem is that the grad of loss wrt the last layer x3 is constant wrt to any of the inputs that require grad, because all you do is a mean().
If you are computing Hessian of loss w.r.t. to parameters, wouldn’t the size be eventually be [B, X, N, N] where X is the numel of the parameter, assuming all parameters have the same numel. Also, if two parameters a, b had different sizes i.e., weight and bias, d2L_dadb would have size of a but d2L_dbda would have size b, so that may confuse things.
From the above though, we seem to be computing Hessian of loss w.r.t. to each of the layer outputs, which is a different quantity.