Some confusion about jacobian_trace

[HoJ-Onle]

Hello. It's a nice work!
I have some confusion about the implementation of the jacobian_trace function. The jacobian_trace function used the function keep_grad which has set create_graph=True. It doesn't seem like we have a need for higher gradient. It had a very high demand on my GPU memory when I was running some larger datasets. I want to know if I can set it to False and only keep retain_graph=True.

def keep_grad(output, input, grad_outputs=None):
    return torch.autograd.grad(output, input, grad_outputs=grad_outputs, retain_graph=True, create_graph=True)[0]

def exact_jacobian_trace(fx, x):
    vals = []
    for i in range(x.size(1)):
        fxi = fx[:, i]
        dfxi_dxi = keep_grad(fxi.sum(), x)[:, i][:, None]
        vals.append(dfxi_dxi)
    vals = torch.cat(vals, dim=1)
    return vals.sum(dim=1)

Looking forward to your relpy.