optax
equinox

--extra-index-url https://download.pytorch.org/whl/cpu
torch==2.0.1+cpu
torchvision==0.15.2+cpu

jax==0.4.28
jaxlib==0.4.28+cuda12.cudnn89 --find-links https://storage.googleapis.com/jax-releases/jax_cuda_releases.html

numpy<2
wandb
tqdm
imageio
opencv-python
scipy
matplotlib
scikit-learn