gistfile1.txt

import torch
from torch.profiler import *

for with_cuda in [True, False]:
    with profile() as prof:
        x = torch.randn(2, 2)
        if with_cuda:
            x = x.cuda()
        x = x.matmul(x)

    print(prof.key_averages().table(sort_by="self_cpu_time_total", row_limit=-1))