gpu_requirement.yml

name: cube
channels:
    - defaults
    - nvidia
    - conda-forge
dependencies:
    - python=3.8
    - pip
    - numpy
    - scipy
    - numba
    - h5py
    - matplotlib
    - scikit-learn
    - jupyter
    - tqdm
    - pytorch-cpu=1.13.0
    - jax=0.3.25
    - jaxlib=0.3.25=*cuda*
    - cudatoolkit=11.7
    - cuda-nvcc=11.7
    - cudnn
    - sentencepiece
    - pip:
        - flax==0.6.0
        - optax==0.1.3
        - chex==0.1.5
        - distrax==0.1.2
        - orbax==0.0.23
        - transformers==4.27.2
        - huggingface_hub==0.13.3
        - datasets==2.9.0
        - einops
        - tensorflow==2.11.0
        - dill
        - seaborn
        - absl-py
        - gtimer
        - joblib
        - wandb==0.13.5
        - ml_collections
        - gcsfs==2022.11.0
        - requests
        - jupyter_http_over_ws
        - lm-eval==0.3.0
        - evaluate==0.3.0
        - pyext==0.7
        - pydantic
        - fastapi
        - uvicorn
        - gradio