(tensor, op=dist.ReduceOp.SUM, world_size=1)
| 123 | |
| 124 | |
| 125 | def all_reduce_tensor(tensor, op=dist.ReduceOp.SUM, world_size=1): |
| 126 | tensor = tensor.clone() |
| 127 | dist.all_reduce(tensor, op) |
| 128 | tensor.div_(world_size) |
| 129 | |
| 130 | return tensor |
| 131 | |
| 132 | |
| 133 | def load_restore_model(model, model_file): |