pytorch>=2.4.1
einops>=0.8.1
