(tensor)
| 65 | |
| 66 | |
| 67 | def all_reduce_sum(tensor): |
| 68 | if not is_distributed(): |
| 69 | return tensor |
| 70 | dim_squeeze = False |
| 71 | if tensor.ndim == 0: |
| 72 | tensor = tensor[None, ...] |
| 73 | dim_squeeze = True |
| 74 | torch.distributed.all_reduce(tensor) |
| 75 | if dim_squeeze: |
| 76 | tensor = tensor.squeeze(0) |
| 77 | return tensor |
| 78 | |
| 79 | |
| 80 | def all_reduce_average(tensor): |
no test coverage detected