Fix GPU-less CPU use on Linux with specific Torch indexes

This commit is contained in:
Luke Stanley 2025-10-21 19:52:21 +00:00
parent defd1246aa
commit 901b075605
3 changed files with 29 additions and 13 deletions

View File

@ -14,7 +14,7 @@ NANOCHAT_BASE_DIR="$HOME/.cache/nanochat"
mkdir -p $NANOCHAT_BASE_DIR mkdir -p $NANOCHAT_BASE_DIR
command -v uv &> /dev/null || curl -LsSf https://astral.sh/uv/install.sh | sh command -v uv &> /dev/null || curl -LsSf https://astral.sh/uv/install.sh | sh
[ -d ".venv" ] || uv venv [ -d ".venv" ] || uv venv
uv sync uv sync --extra cpu
source .venv/bin/activate source .venv/bin/activate
if [ -z "$WANDB_RUN" ]; then if [ -z "$WANDB_RUN" ]; then
WANDB_RUN=dummy WANDB_RUN=dummy

View File

@ -44,11 +44,11 @@ python_files = ["test_*.py"]
python_classes = ["Test*"] python_classes = ["Test*"]
python_functions = ["test_*"] python_functions = ["test_*"]
# target torch to cuda 12.8 # target torch to cuda 12.8 or CPU
[tool.uv.sources] [tool.uv.sources]
torch = [ torch = [
{ index = "pytorch-cpu", marker = "sys_platform != 'linux'" }, { index = "pytorch-cpu", extra = "cpu" },
{ index = "pytorch-cu128", marker = "sys_platform == 'linux'" }, { index = "pytorch-cu128", extra = "gpu" },
] ]
[[tool.uv.index]] [[tool.uv.index]]
@ -60,3 +60,19 @@ explicit = true
name = "pytorch-cu128" name = "pytorch-cu128"
url = "https://download.pytorch.org/whl/cu128" url = "https://download.pytorch.org/whl/cu128"
explicit = true explicit = true
[project.optional-dependencies]
cpu = [
"torch>=2.8.0",
]
gpu = [
"torch>=2.8.0",
]
[tool.uv]
conflicts = [
[
{ extra = "cpu" },
{ extra = "gpu" },
],
]

View File

@ -23,7 +23,7 @@ command -v uv &> /dev/null || curl -LsSf https://astral.sh/uv/install.sh | sh
# create a .venv local virtual environment (if it doesn't exist) # create a .venv local virtual environment (if it doesn't exist)
[ -d ".venv" ] || uv venv [ -d ".venv" ] || uv venv
# install the repo dependencies # install the repo dependencies
uv sync uv sync --extra gpu
# activate venv so that `python` uses the project's venv instead of system python # activate venv so that `python` uses the project's venv instead of system python
source .venv/bin/activate source .venv/bin/activate