torchvision
matplotlib
optax