(x)
| 330 | |
| 331 | |
| 332 | def all_reduce_mean(x): |
| 333 | world_size = get_world_size() |
| 334 | if world_size > 1: |
| 335 | x_reduce = torch.tensor(x).cuda() |
| 336 | dist.all_reduce(x_reduce) |
| 337 | x_reduce /= world_size |
| 338 | return x_reduce.item() |
| 339 | else: |
| 340 | return x |
nothing calls this directly
no test coverage detected