This puzzles me a bit too - new codegen creates over-specialized kernels, hardcoding tensor sizes. However, my understanding is, it will compile multiple times for different sizes, if it sees the same shape[s] twice. There is torch._C._jit_set_bailout_depth(n) option to control this, it is 20 by default.
Simpler alternative is to use legacy jit, it produces more generic code.