Does VICREG loss work correctly in multi-gpus setting?

[NoTody]

Hi, Kevin

I notice that the current implementation of VICREG doesn't have the gather layer for the input embeddings. Does it still work as expected to get embeddings from mini-batches in multi-gpus in this case? See https://github.com/facebookresearch/vicreg/blob/a73f567660ae507b0667c68f685945ae6e2f62c3/main_vicreg.py#L200 for original implementation of VICREG.

Regards,
Notody

[KevinMusgrave]

Yeah I don't think it does. For me, the ideal solution would be to make it compatible with DistributedLossWrapper. That means either

• making it have the same input format as BaseMetricLossFunction. I don't know if there's a way to do this that makes sense.
• Adding an if-statement in DistributedLossWrapper to catch the VICReg loss type

[KevinMusgrave]

In v2.0 this will work:

loss_fn = DistributedLossWrapper(loss=VICRegLoss())
loss = loss_fn(embeddings, ref_emb=ref_emb)