Find a way to compile language models

[Francesco215]

torch.compile doesn't seem to work with language models because the input sequence is always different