Opacus only captures the first argument to a module's forward as activations in the capture_activations_hook and misses out on other args such as offsets in case of nn.EmbeddingBag
This results in an error during per-sample grad computation even with functorch.
alexandresablayrolles
commented
1 year ago
🐛 Bug
Opacus only captures the first argument to a module's forward as activations in the
capture_activations_hookand misses out on other args such asoffsetsin case ofnn.EmbeddingBagThis results in an error during per-sample grad computation even with functorch.
COLAB: https://drive.google.com/file/d/1wvfO3roibmYxtHV_SBcdW3PjkrkCxalW/view?usp=sharing