Wrapping the loss with DistributedDataParallel

[elisim]

Hi,

When using DDP, should the loss_fn be wrapped with DistributedDataParallel? I’m specifically working with CosFace and ArcFace, both of which have a W parameter in the loss function. To ensure the gradients of W are synchronized across all processes, is it necessary to wrap loss_fn with DistributedDataParallel?

I saw it here:

#218

[KevinMusgrave]

Yeah I think it's necessary.