Hi @Torcione
Yes, you can use opacus for that; Take a look at GradSampleModule (opacus/grad_sample/grad_sample_module.py)
It’s a wrapper around nn.Module that encapsulates per sample gradient computation. When you wrap your model with GradSampleModule, each trainable parameter will get .grad_sample attribute containing per-sample gradients. No noise or clipping is performed.
Hope this helps