jax
optax
torch
torchvision
torchmetrics
matplotlib
