Gather and stack/cat values from all processes, if there are multiple processes.
(values: torch.Tensor, rank: int, world_size: int)
| 81 | |
| 82 | |
| 83 | def all_gather_if_needed(values: torch.Tensor, rank: int, world_size: int) -> torch.Tensor: |
| 84 | """Gather and stack/cat values from all processes, if there are multiple processes.""" |
| 85 | if world_size == 1: |
| 86 | return values |
| 87 | |
| 88 | all_values = [torch.empty_like(values).to(rank) for _ in range(world_size)] |
| 89 | dist.all_gather(all_values, values) |
| 90 | cat_function = torch.cat if values.dim() > 0 else torch.stack |
| 91 | return cat_function(all_values, dim=0) |
| 92 | |
| 93 | |
| 94 | def formatted_dict(d: Dict) -> Dict: |
no outgoing calls
no test coverage detected