How to get step time from torch.profiler?

1

I would like to get time for a step from torch profiler. Is there any way to do?

skip_first = 5
wait = 5
warmup = 5
active = 5
repeat = 3

num_steps = wait + warmup + active
iterations = skip_first + num_steps * repeat

with torch.profiler.profile(
            activities=[
                torch.profiler.ProfilerActivity.CPU,
                torch.profiler.ProfilerActivity.CUDA,
            ],
            schedule=torch.profiler.schedule(
                skip_first=skip_first,
                wait=wait,
                warmup=warmup,
                active=active,
                repeat=repeat,
            ),
            on_trace_ready=torch.profiler.tensorboard_trace_handler("./log"),
        ) as p:
            for iter in range(iterations):
                model(input)
                p.step()
print(p.key_averages().total_average().self_cuda_time_total)