Skip to content

fix(rerank): pin torch to cu128 wheels for wider CUDA compatibility#152

Closed
shan-nvidia wants to merge 1 commit intooholworthy/rerank-recipe-v1from
steve/rerank_recipe_patch1
Closed

fix(rerank): pin torch to cu128 wheels for wider CUDA compatibility#152
shan-nvidia wants to merge 1 commit intooholworthy/rerank-recipe-v1from
steve/rerank_recipe_patch1

Conversation

@shan-nvidia
Copy link
Copy Markdown
Contributor

@shan-nvidia shan-nvidia commented Apr 16, 2026

Problem

The rerank eval stage (stage3_eval) resolved PyTorch from PyPI, which
pulls wheels compiled against the latest CUDA toolkit (currently 13.0).
This breaks on any server whose NVIDIA driver does not yet support
CUDA 13.0, for example:

CUDA initialization: The NVIDIA driver on your system is too old (found version 12090).

Fix

Pin torch to the cu128 (CUDA 12.8) wheel index, which is
forward-compatible with CUDA 12.8, 12.9, 13.0, and newer drivers.

  • Added torch as a direct dependency (previously only a transitive dep
    via sentence-transformers).
  • Configured a [[tool.uv.index]] pointing to
    https://download.pytorch.org/whl/cu128 with explicit = true, so
    only torch resolves from it — all other packages continue to use
    PyPI.
  • Added [tool.uv.sources] mapping torch to the cu128 index.
  • Regenerated uv.lock.

Files changed

File Change
src/nemotron/recipes/rerank/stage3_eval/pyproject.toml torch cu128 index + direct dep
src/nemotron/recipes/rerank/stage3_eval/uv.lock Regenerated

Test plan

  • Run nemotron rerank eval on a CUDA 12.x driver — verify torch initializes CUDA successfully

Signed-off-by: Steve Han <sthan@nvidia.com>
@shan-nvidia shan-nvidia self-assigned this Apr 16, 2026
@shan-nvidia
Copy link
Copy Markdown
Contributor Author

Close this since it's already fixed.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant