distinguish CPU and GPU setups

.github/workflows/python-package-conda.yml

@@ -16,10 +16,10 @@ jobs:
         uses: actions/setup-python@v3
         with:
           python-version: ${{ matrix.python-version }}
-      - name: Install Conda environment from `environment.yml` with Micromamba
+      - name: Install Conda environment from `environment-cpu.yml` with Micromamba
         uses: mamba-org/setup-micromamba@v1
         with:
-          environment-file: environment.yml
+          environment-file: environment-cpu.yml
           init-shell: >-
             bash
             zsh

environment-cpu.yml

@@ -0,0 +1,28 @@
+name: TODO_package_name
+
+channels:
+  - default
+  - conda-forge
+
+dependencies:
+  - python>=3.11
+  - ipython
+  - ipykernel
+  - pre-commit
+
+  - numpy
+  - pandas>=2.0
+  - pytables
+  - submitit
+
+  # PyTorch:
+  # - pytorch::pytorch
+  # - pytorch::torchvision
+  # - huggingface::transformers
+
+  # JAX:
+  # - conda-forge::jax
+
+  - pip
+  - pip:
+      - -e .

environment-gpu.yml

@@ -0,0 +1,37 @@
+name: TODO_package_name
+
+channels:
+  - default
+  - conda-forge
+
+dependencies:
+  - python>=3.11
+  - ipython
+  - ipykernel
+  - pre-commit
+
+  - numpy
+  - pandas>=2.0.0
+  - pytables
+  - submitit
+
+  # NOTE: The following requires a GPU node available
+  # at the time of creating the environment.
+  # CUDA:
+  # - conda-forge::cudatoolkit>=11.8
+  # - nvidia::cuda>=11.8
+
+  # PyTorch:
+  # - pytorch::pytorch-cuda>=12.0
+  # - pytorch::pytorch
+  # - pytorch::torchvision
+  # - huggingface::transformers
+
+  # JAX:
+  # - nvidia::cuda-nvcc>=12.0
+  # - conda-forge::jaxlib=*=*cuda*
+  # - conda-forge::jax
+
+  - pip
+  - pip:
+      - -e .