Hi,
If you don’t run inside DataParallel, this will work fine.
The thing is that DataParallel has a lot of logic on how to get the different copies on different devices and how to sync the gradients.
In particular I think that it will accumulate the gradients of your subnet automatically.
You can check this by using simple functions with known gradients. (Like addition that just backwards the original gradient)