Hey
Sorry, but I’m still having trouble reproducing the issue.
I’ve tried multiple package versions (opacus 0.13, 0.14, master), but none produce the error you’ve described.
Can you maybe share a Colab notebook with the error to help find the reason?
PS: While investigating this we’ve found and fixed 6 quite bad memory inefficiency, so thanks for pointing that way