torch
transformers
pytorch-lightning
einops
