test_dynamic.py

import torch
from torch._inductor import config

batch_size = 4
in_features = 512
out_features = 1024
dtype = torch.bfloat16
# dtype = torch.float16
bias = True


class M(torch.nn.Module):
    def __init__(self, bias):
        super().__init__()
        self.linear0 = torch.nn.Linear(in_features, out_features, bias=False)

    def forward(self, x):
        return torch.relu(x)

if __name__ == "__main__":
    with torch.no_grad():
        input = torch.randn(batch_size, in_features, dtype=dtype)

        m = M(bias=bias).eval()
        ref_res = m(input)
        cm = torch.compile(m)
        act_res = cm(input)

        input2 = torch.randn(batch_size+2, in_features, dtype=dtype)
        act_res = cm(input2)