Upgrade to torch 2.0

[agastyaseth]

are there any plans to port the library to torch 2? Since the parallelize() library is deprecated in torch 2, it becomes impossible to train larger models like llama 7b and mistral 7b even with A100 80GBs