Currently I use the following. Is there a better way to enable it without manually calling __enter__? Is it necessary (I came up with it when it seemed necessary, but now it was maybe refactored?)?
if args.profile_autograd:
    autograd_profiler = torch.autograd.profiler.profile()
    autograd_profiler.__enter__()

# model running

if args.profile_autograd:
    autograd_profiler.__exit__(None, None, None)
    autograd_profiler.export_chrome_trace(args.profile_autograd)

I don’t want to use with construct because I want to keep enabling the profiler under the flag and prefer not to factor out the model code in a separate function. enable()-kind of API exists for autograd itself, so I thought maybe it exists for the profiler as well. It also exists for nvprof: torch.cuda.profiler.start()