Hi @Yuancheng_Xu!
The hook for computing grad_sample doesn’t work correctly with more than one backward pass.
I’m not sure I completely understand your use case, but I think you can do the following:
# initialize the vanilla model
model = Model(10,num_classes)
model.zero_grad()
X = torch.rand(bs,10)
X.requires_grad = True
grad_x = torch.zeros(bs,num_classes,10)

# compute the output and doutput/dinput
output = model(X) # bs * num_classes
for c in range(num_classes):
    grad_x[:,c,:] = grad(
        outputs=output[:,c],
        inputs=X,
        grad_outputs=torch.ones_like(output[:,c]),
        retain_graph=True
    )[0]

# add the hook & evaluate the model one more time to get the grad_sample
dp_model = GradSampleModule(model)
output = dp_model(X)

# ... apply gradient step with grad_sample

# remove hooks before the next iteration
# if you don't remove hooks, model(data) will be computed with grad_samples leading to the exception
# also you won't be able to recreate GradSampleModule(model) before removing hooks
dp_model.remove_hooks() 

Note that using grads without properly added noise like in this case could break DP guarantees if you expected any.