Code explanation in gather.py

[littleredxh]

Hi, Janne

The GatherLayer module in gather.py is smart and efficient. I really appreciate this module.

I almost understand most of the functions of the codes but still confuse at one point.

line 19 of SimCLR/simclr/modules/gather.py

    grad_out[:] = grads[dist.get_rank()]

Would you mind explaining it in more detail?

[ivallesp]

Basically, the gradients in each device would be the local gradients. Look at this blog post: https://amsword.medium.com/gradient-backpropagation-with-torch-distributed-all-gather-9f3941a381f8