torch
torchvision
matplotlib
numpy
wandb