gather data from multiple processes to one list
(data, return_np=True)
| 9 | |
| 10 | |
| 11 | def gather_data(data, return_np=True): |
| 12 | ''' gather data from multiple processes to one list ''' |
| 13 | data_list = [torch.zeros_like(data) for _ in range(dist.get_world_size())] |
| 14 | dist.all_gather(data_list, data) # gather not supported with NCCL |
| 15 | if return_np: |
| 16 | data_list = [data.cpu().numpy() for data in data_list] |
| 17 | return data_list |
| 18 | |
| 19 | def autocast(f): |
| 20 | def do_autocast(*args, **kwargs): |
nothing calls this directly
no outgoing calls
no test coverage detected