Performs all_gather operation on the provided tensors. *** Warning ***: torch.distributed.all_gather has no gradient.
(tensor)
| 383 | # utils |
| 384 | @torch.no_grad() |
| 385 | def concat_all_gather(tensor): |
| 386 | """ |
| 387 | Performs all_gather operation on the provided tensors. |
| 388 | *** Warning ***: torch.distributed.all_gather has no gradient. |
| 389 | """ |
| 390 | tensors_gather = [ |
| 391 | torch.ones_like(tensor) |
| 392 | for _ in range(torch.distributed.get_world_size()) |
| 393 | ] |
| 394 | torch.distributed.all_gather(tensors_gather, tensor, async_op=False) |
| 395 | |
| 396 | output = torch.cat(tensors_gather, dim=1) |
| 397 | return output |
| 398 | |
| 399 | |
| 400 | class DINOLoss(nn.Module): |
nothing calls this directly
no outgoing calls
no test coverage detected