use_profiler.py

import torch

def fn(x, y):
    return torch.cat([x + y, y]).sin()

a = torch.ones((1024, 256), dtype=torch.float32)
b = torch.ones((1024, 256), dtype=torch.float32) * 2

with torch.profiler.profile(schedule=torch.profiler.schedule(wait=2, warmup=2, repeat=1, active=2), record_shapes=True) as prof:
    for _ in range(8):
        fn(a, b)
        prof.step()

prof.export_chrome_trace("/home/dberard/matmul_prof.json")