env_gpu.yml

name: tsadar-gpu
channels:
  - conda-forge
dependencies:
  - python=3.11
  - pip
  - pyhdf
  - xlrd
  - pip:
      - "jax[cuda12]"
      - jaxopt
      - numpy
      - scipy
      - matplotlib
      - pyyaml
      - mlflow
      - boto3
      - flatten-dict
      - typing-extensions
      - optax
      - tqdm
      - xarray
      - mlflow_export_import
      - pandas
      - interpax
      - tabulate
      - psutil
      - pynvml