An all_gather layer with backward, useful for collecting model output embeddings from multiple gpus to allow large batch size loss calculation, e.g. for InfoNCE (SimCRL, CLIP).

class GatherLayer[source]

GatherLayer(*args, **kwargs) :: Function

Gather tensors from all process, supporting backward propagation. https://github.com/open-mmlab/OpenSelfSup/blob/696d04950e55d504cf33bc83cfadbb4ece10fbae/openselfsup/models/utils/gather_layer.py