Use CPU/GPU extras with uv sources for torch installation

.github/workflows/install-and-test.yml

@@ -1,23 +1,25 @@
 # Configuration to test both CPU and GPU installations with both UV and pip
 #
-# We don't use lock files in neural-lam, because that would tie us to a
-# particular variant of pytorch (gpu/cpu), and want it to be possible for
-# people choose which variant to install. Because we don't use lock files the
-# specific package versions to install is resolved on every install. This means
-# that has upstream packages evolve the specific version installed might change
-# (if we don't pin it). Currently, we do not pin pytorch in pyproject.toml
+# torch is a required dependency in [project.dependencies]. The CPU/GPU variant
+# is selected at install time via extras defined in pyproject.toml:
+#
+#   uv sync --extra cpu --group dev  # CPU runners
+#   uv sync --extra gpu --group dev  # GPU runners (CUDA 12.8)
+#
+# The extras route torch to the correct PyTorch index via [tool.uv.sources]
+# and [[tool.uv.index]] in pyproject.toml.
+#
+# NOTE: `uv sync --torch-backend` is not yet supported (uv pip only).
+# Tracking: https://github.com/astral-sh/uv/issues/12994
+#
+# pip does not support extras-based index routing. The pip path uses a dry-run
+# to detect the required torch version and installs from the correct index URL.
 #
 # To test both CPU and GPU installations with both UV and pip we use two sets
 # of hosts to run the training on:
 # - ubuntu-latest for CPU: these are on github
 # - cirun-aws-runner--${{ github.run_id }} for GPU: these are on AWS,
 #   orchestrated by cirun.io
-#
-# To enable installing different pytorch variants we use an approach where we
-# first install a specific torch version (and variant), and then install neural-lam.
-# To ensure that the version of torch installed is compatible with neural-lam we
-# first use `pip install --dry-run .` to determine the version of torch that would
-# be installed.
 name: CPU+GPU testing
 
 on: [push, pull_request, workflow_dispatch]
@@ -48,7 +50,7 @@ jobs:
         with:
           python-version: 3.13
 
-      - name : Install uv (if applicable)
+      - name: Install uv (if applicable)
         if: matrix.package_manager == 'uv'
         uses: astral-sh/setup-uv@v7
 
@@ -66,24 +68,20 @@ jobs:
           echo "Torch version detected: $TORCH_VERSION"
           echo "TORCH_VERSION=$TORCH_VERSION" >> $GITHUB_ENV
 
-      - name: Install PyTorch
-        if : matrix.package_manager != 'uv'
+      - name: Install PyTorch (pip)
+        if: matrix.package_manager == 'pip'
         run: |
           python -m pip install "torch==$TORCH_VERSION" --index-url ${{ matrix.torch_index }}
 
-      - name: Install package (including dev dependencies)
-        if : matrix.package_manager != 'uv'
+      - name: Install package (pip)
+        if: matrix.package_manager == 'pip'
         run: |
           python -m pip install --group dev .
 
-      - name: Setup environment with uv (if applicable)
+      - name: Install with uv
         if: matrix.package_manager == 'uv'
         run: |
-          uv venv --no-project
-          source .venv/bin/activate
-          uv pip install torch --index-url ${{ matrix.torch_index }}
-          uv pip install --group dev -e .
-          uv pip list
+          uv sync --extra ${{ matrix.device }} --group dev --locked
           echo "$PWD/.venv/bin" >> $GITHUB_PATH
 
       - name: Print torch version

CHANGELOG.md

@@ -23,6 +23,9 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 - Add comprehensive type hints to `neural_lam/create_graph.py` [\#618](https://github.com/mllam/neural-lam/pull/618) @GiGiKoneti
 
+- Use `cpu`/`gpu` extras with `[tool.uv.sources]` and a committed `uv.lock` to simplify and pin torch installation in CI and for users [\#604](https://github.com/mllam/neural-lam/pull/604) @RajdeepKushwaha5
+- Update default CUDA version in install instructions and CI from 11.1 to 12.8. Users on older CUDA versions need to install the matching `torch` variant manually (see README) [\#604](https://github.com/mllam/neural-lam/pull/604) @RajdeepKushwaha5
+
 ## [v0.6.0](https://github.com/mllam/neural-lam/releases/tag/v0.6.0)
 
 This release introduces new features including GIF animation support, wandb run resumption, and improved ensemble loading, alongside a large number of bug fixes and maintenance updates.

README.md

@@ -70,10 +70,7 @@ keeping versions consistent (it automatically updates the `pyproject.toml`
 file), makes it easy to handle virtual environments and includes the
 development toolchain packages installation too.
 
-**regarding `torch` installation**: because `torch` creates different package
-variants for different CUDA versions and cpu-only support you will need to install
-`torch` separately if you don't want the most recent GPU variant that also
-expects the most recent version of CUDA on your system.
+**regarding `torch` installation**: `torch` is a required dependency, but it ships separate variants for CPU-only and each CUDA version. With `uv`, the correct PyTorch index is chosen automatically via `[tool.uv.sources]`. With `pip`, you must pass `--index-url` explicitly (see instructions below).
 
 We cover all the installation options in our [github actions ci/cd
 setup](.github/workflows/) which you can use as a reference.
@@ -90,16 +87,27 @@ python -m pip install neural_lam
 
 1. Clone this repository and navigate to the root directory.
 2. Install `uv` if you don't have it installed on your system (either with `pip install uv` or [following the install instructions](https://docs.astral.sh/uv/getting-started/installation)).
-> If you are happy using the latest version of `torch` with GPU support (expecting the latest version of CUDA is installed on your system) you can skip to step 5.
-3. Create a virtual environment for uv to use with `uv venv --no-project`.
-4. Install a specific version of `torch` with `uv pip install torch --index-url https://download.pytorch.org/whl/cpu` for a CPU-only version or `uv pip install torch --index-url https://download.pytorch.org/whl/cu111` for CUDA 11.1 support (you can find the correct URL for the variant you want on [PyTorch webpage](https://pytorch.org/get-started/locally/)).
-5. Install the dependencies with `uv pip install .`. If you will be developing `neural-lam` we recommend to install the development dependencies with `uv pip install --group dev -e .`. This installs the `neural-lam` package in editable mode, so you can make changes to the code and see the effects immediately.
+3. Install with the CPU or GPU extra:
+   ```bash
+   uv sync --extra cpu --group dev  # CPU-only
+   uv sync --extra gpu --group dev  # GPU (CUDA 12.8)
+   ```
+   This creates a virtual environment, installs `torch` from the correct
+   PyTorch index, and installs all other dependencies (including dev tools)
+   in one step.
+
+   If you need a different CUDA version, install with the `cpu` extra first (smaller download) and then swap `torch` to your CUDA variant:
+   ```bash
+   uv sync --extra cpu --group dev              # create env (small torch wheel)
+   uv pip install torch --torch-backend auto    # replace torch with your CUDA variant
+   ```
+   Note: this swap is venv-local and will be reverted by the next `uv sync`.
 
 #### Using `pip`
 
 1. Clone this repository and navigate to the root directory.
 > If you are happy using the latest version of `torch` with GPU support (expecting the latest version of CUDA is installed on your system) you can skip to step 3.
-2. Install a specific version of `torch` with `python -m pip install torch --index-url https://download.pytorch.org/whl/cpu` for a CPU-only version or `python -m pip install torch --index-url https://download.pytorch.org/whl/cu111` for CUDA 11.1 support (you can find the correct URL for the variant you want on [PyTorch webpage](https://pytorch.org/get-started/locally/)).
+2. Install a specific version of `torch` with `python -m pip install torch --index-url https://download.pytorch.org/whl/cpu` for a CPU-only version or `python -m pip install torch --index-url https://download.pytorch.org/whl/cu128` for CUDA 12.8 support (you can find the correct URL for the variant you want on [PyTorch webpage](https://pytorch.org/get-started/locally/)).
 3. Install the dependencies with `python -m pip install .`. If you will be developing `neural-lam` we recommend to install in editable mode and install the development dependencies with `python -m pip install --group dev -e .` so you can make changes to the code and see the effects immediately.
 
 

pyproject.toml

@@ -41,9 +41,36 @@ dependencies = [
 ]
 requires-python = ">=3.10"
 
+[project.optional-dependencies]
+cpu = ["torch>=2.3.0"]
+gpu = ["torch>=2.3.0"]
+
 [dependency-groups]
 dev = ["pre-commit>=3.8.0", "pytest>=8.3.2", "pooch>=1.8.2"]
 
+[tool.uv]
+# uv-specific configuration for PyTorch CPU/GPU variant selection.
+# NOTE: --torch-backend is only supported in `uv pip`, not `uv sync`.
+# Tracking issue for project-interface support: https://github.com/astral-sh/uv/issues/12994
+# Until then, we use extras + explicit index routing for `uv sync`.
+conflicts = [[{ extra = "cpu" }, { extra = "gpu" }]]
+
+[[tool.uv.index]]
+name = "pytorch-cpu"
+url = "https://download.pytorch.org/whl/cpu"
+explicit = true
+
+[[tool.uv.index]]
+name = "pytorch-gpu"
+url = "https://download.pytorch.org/whl/cu128"
+explicit = true
+
+[tool.uv.sources]
+torch = [
+    { index = "pytorch-cpu", extra = "cpu" },
+    { index = "pytorch-gpu", extra = "gpu" },
+]
+
 [tool.black]
 line-length = 80