torch
torchvision
einops