def profiler_runner()

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