dl2023_jax_gpu.yml

name: dl2023_jax
channels:
  - pytorch
  - nvidia
  - conda-forge
  - defaults
dependencies:
  - python=3.10.6
  - pip=22.2.2
  - pytorch-cuda=11.7
  - pytorch=1.13.0
  - jax=0.4.5
  - jaxlib[build=*cuda*]
  - cuda-nvcc
  - torchvision=0.14.0
  - torchaudio=0.13.0
  - pip:
    - pytorch-lightning==1.7.7
    - tensorboard==2.11.2
    - tensorflow==2.11.0
    - optax==0.1.4
    - flax==0.6.6
    - tabulate>=0.8.9
    - tqdm>=4.62.3
    - pillow>=8.0.1
    - notebook>=6.4.5
    - jupyterlab>=3.2.1
    - matplotlib>=3.4.3
    - seaborn>=0.11.2
    - ipywidgets>=7.6.5