Sync all GPUs to make sure all operations are finished, needed for correct benchmarking of latency/throughput.
()
| 183 | return loader |
| 184 | |
| 185 | def sync_gpus() -> None: |
| 186 | """Sync all GPUs to make sure all operations are finished, needed for correct benchmarking of latency/throughput.""" |
| 187 | for i in range(torch.cuda.device_count()): |
| 188 | torch.cuda.synchronize(device=i) |
| 189 | |
| 190 | def map_tensors(obj, device: torch.device | str | None = None, dtype: torch.dtype | None = None): |
| 191 | """Recursively map tensors to device and dtype.""" |