torch
torchvision
Pillow
einops
numpy