(ctx, tensor)
| 458 | |
| 459 | @staticmethod |
| 460 | def forward(ctx, tensor): |
| 461 | ctx.batch_size = tensor.shape[0] |
| 462 | |
| 463 | gathered_tensor = [torch.zeros_like(tensor) for _ in range(torch.distributed.get_world_size())] |
| 464 | |
| 465 | torch.distributed.all_gather(gathered_tensor, tensor) |
| 466 | gathered_tensor = torch.cat(gathered_tensor, 0) |
| 467 | |
| 468 | return gathered_tensor |
| 469 | |
| 470 | @staticmethod |
| 471 | def backward(ctx, grad_output): |
nothing calls this directly
no outgoing calls
no test coverage detected