ِDo you mean I need to create a function like: (W: weight matrix between input x and h layer, b: bias term)
def func(x):
    fc_out=torch.mm(x,W)+b
    return fc_out

jacobian(func,x)

shall I set “create_graph=True”?