torch
torchvision
einops
matplotlib
numpy
