randophilus:

While this works in the network by itself, it only returns the original value of 0 when I wrap it with dataparallel.


I believe this is because DataParallel (DP) would replicate the module to each provided device in every forward. So the module instance stored in DataParallel and the module instance used by forward on every device are different.
However, even if they are the same instance, it’s still problematic. Suppose we run DP on two devices, then, one of the device would override the value of unsupervised_loss written by another device. Is this an expected behavior?

I instead rewrite the function to return both the conv output and the unsupervised loss as a tuple.

This seems a reasonable workaround.