I usually prefer torch.autograd.grad 23 for anything except differentiating a target function to take an optimizer step.
That also keeps you from accumulating things the wrong things…
So instead of dfdx(x, z, self.functional) you do
f = self.functional(x,z) and
dfdx,dfdz = torch.autograd.grad([f],[x,z], retain_graph=True, create_graph=True). (retain_graph only if you have need for it, i.e. get called backward twice errors without.)
Best regards
Thomas