torch>=1.8
einops==0.3.2
