Grad Sample Module: Use full backward hook to save activations and backprop values.

🐛 Bug

I am training my ConvNet model with OCT Data and analysing the privacy spent using Opacus by implementing the Random Sparsification which uses backPACK for calculating gradients.

Please reproduce using this repo: Github

To Reproduce

:warning: We cannot help you without you sharing reproducible code. Do not ignore this part :) Steps to reproduce the behavior:

python code/train_with_rs_opacus.py

Using a non-full backward hook when the forward contains multiple autograd Nodes is deprecated and will be removed in future versions. This hook will be missing some grad_input. Please use register_full_backward_hook to get the documented behavior.

Expected behavior

The Grad sample module should use full backward hook instead of regular hook.

self.autograd_grad_sample_hooks.append(
    module.register_full_backward_hook(
        partial(
            self.capture_backprops_hook,
            loss_reduction=loss_reduction,
            batch_first=batch_first,
        )
    )
)

Environment

PyTorch Version (e.g., 1.0): 2.3.0
OS (e.g., Linux): MacOS M2
How you installed PyTorch (conda, pip, source): pip
Build command you used (if compiling from source):
Python version: 3.11.8
CUDA/cuDNN version:
GPU models and configuration:
Any other relevant information: backpack-for-pytorch==1.6.0
Opacus Version: 1.4.1

Additional context

I would be happy to contribute and implement this, need a bit of direction on what should be changed. Thanks :)

pytorch / opacus