import torch
from rpdTracerControl import rpdTracerControl
for step in [1,2,3]:
trace_file_name = f"{step}_torch.rpd"
print(f"in {step}!!")
with rpdTracerControl(trace_file_name, nvtx=True) as p:
a =torch.rand(step).cuda()
torch.add(a, 20)
p.top_totals()
[x] Make rpdTracerControl into a standalone context manager similar to the functionality of torch.profile
[x] Make it all callable in a loop, regardless of the static nature of the library
The goal is to allow something like: