BatchMemoryManager for training private GANs

[Caarinaaa]

🚀 Feature

BatchMemoryManager support for multiple optimizers with some of which are DPOptimizer.

Motivation

Hi! I am trying to implement BatchMemoryManager when training GAN.
It seems that BatchMemoryManager need the optimizer, but I got two optimizers here. One optimizerG for generator and the other optimizerD for discriminator. Moreover, optimizerD is made private by Opacus while optimizerG is not. So how can I use BatchMemoryManager here?

Pitch

Alternatives

Additional context