| 163 | |
| 164 | |
| 165 | class MeasureTime(): |
| 166 | def __init__(self, measurements, key, cpu_run=False): |
| 167 | self.measurements = measurements |
| 168 | self.key = key |
| 169 | self.cpu_run = cpu_run |
| 170 | |
| 171 | def __enter__(self): |
| 172 | if not self.cpu_run: |
| 173 | torch.cuda.synchronize() |
| 174 | self.t0 = time.perf_counter() |
| 175 | |
| 176 | def __exit__(self, exc_type, exc_value, exc_traceback): |
| 177 | if not self.cpu_run: |
| 178 | torch.cuda.synchronize() |
| 179 | self.measurements[self.key] = time.perf_counter() - self.t0 |
| 180 | |
| 181 | |
| 182 | def main(): |
no outgoing calls