[FEATURE-REQUEST] Switch to Pytorch 2.0, and use compiled models.

[AJDERS]

Pytorch 2.0 has been released, and promises backwards compatibility, and many performance enhancements.

As an example of the performance enhancements, an option to compile models is now available, resulting in a ~50% speed-up on huggingface models on an A100, and slightly less on non-server-class GPUs, with minimal code. The speed-ups are even large for AMP precision compared to float32.

Seems like simply doing:

compiled_model = torch.compile(model)

will lead to this speed-up.

[AJDERS]

This would constitute a major release, and is seen as nice to have.