in run_profile.py [0:0]
def profiler_runner(path, fn, *args, **kwargs):
with torch.profiler.profile(
activities=[torch.profiler.ProfilerActivity.CPU, torch.profiler.ProfilerActivity.CUDA], record_shapes=True
) as prof:
result = fn(*args, **kwargs)
prof.export_chrome_trace(path)
return result