import torch
from torch import nn
from torch.profiler import ProfilerActivity, profile, record_function


torch.manual_seed(7)

model = nn.Sequential(
    nn.Linear(8, 16),
    nn.ReLU(),
    nn.Linear(16, 4),
)
model.eval()
inputs = torch.randn(32, 8)

with torch.inference_mode():
    for _ in range(5):
        model(inputs)

with torch.inference_mode():
    with profile(
        activities=[ProfilerActivity.CPU],
        record_shapes=True,
    ) as profiler:
        with record_function("model_inference"):
            for _ in range(10):
                model(inputs)

print(
    profiler.key_averages().table(
        sort_by="self_cpu_time_total",
        row_limit=6,
    )
)
profiler.export_chrome_trace("profile_trace.json")