--extra-index-url https://download.pytorch.org/whl/cu118
torch==2.2.0+cu118
torchvision==0.17.0
numpy==1.26.4
wandb
einops
tqdm
