Opacus ValueError: Per sample gradient is not initialized. Not updated in backward pass?

Hello,

I have recently started working on Machine Learning and Medical Imaging and am currently attempting to integrate Opacus into the “Robbing the Fed” attack from the Breaching framework for federated learning. Unfortunately, I encountered an error that I haven’t been able to resolve.
When I initiate training, the following error occurs:

ValueError: Per sample gradient is not initialized. Not updated in backward pass?

The error appears in the file training_opacus.ipynb in my repository:

As I am still a novice programmer, I would appreciate it if you could help me understand the cause of this error and how to fix it. Also, how can I remove the opacus wrapper after training to just keep the model for evaluation?

Kind regards.