accelerate
torch
torchvision
datasets
scipy
wandb
einops
fire
retry
kornia==0.5.4