Obtain the mean of tensor on different GPUs.
(tensor)
| 59 | |
| 60 | |
| 61 | def reduce_mean(tensor): |
| 62 | """"Obtain the mean of tensor on different GPUs.""" |
| 63 | if not (dist.is_available() and dist.is_initialized()): |
| 64 | return tensor |
| 65 | tensor = tensor.clone() |
| 66 | dist.all_reduce(tensor.div_(dist.get_world_size()), op=dist.ReduceOp.SUM) |
| 67 | return tensor |