--extra-index-url https://download.pytorch.org/whl/cu121

# cu121 is to meet our nvidia env. you can change it to other versions
torch
torchvision

timm
diffusers
accelerate
torchdiffeq
pytorch_fid
tensorboard==2.16.2
omegaconf==2.3.0
einops
fairscale
safetensors

# # use for reconstruction evaluation, optional
torchmetrics 
