(func, *args, **kwargs)
| 355 | |
| 356 | |
| 357 | def benchmark(func, *args, **kwargs): |
| 358 | torch.cuda.synchronize() |
| 359 | torch.cuda.reset_peak_memory_stats() |
| 360 | begin = timeit.default_timer() |
| 361 | output = func(*args, **kwargs) |
| 362 | torch.cuda.synchronize() |
| 363 | end = timeit.default_timer() |
| 364 | return output, (end - begin), torch.cuda.max_memory_allocated() |
| 365 | |
| 366 | |
| 367 | ############################################################################## |
no test coverage detected