torch>=1.6
einops>=0.6.1
datasets
accelerate
transformers
optimus-prime-transformers
lion_pytorch
