.devcontainer/download_libtorch.sh (1)

7-7: Make libtorch version configurable to reduce future churn

Hard-coding the version forces PRs for every bump. Allow overriding via env var while defaulting to 2.8.0+cpu.

Apply this diff to the download line:

-wget https://download.pytorch.org/libtorch/cpu/libtorch-cxx11-abi-shared-with-deps-2.8.0%2Bcpu.zip -O ~/libtorch.zip
+wget "https://download.pytorch.org/libtorch/cpu/libtorch-cxx11-abi-shared-with-deps-${LIBTORCH_VERSION:-2.8.0+cpu}.zip" -O ~/libtorch.zip

And add this near the top (outside the changed hunk) to document usage:

# Optional: export LIBTORCH_VERSION=2.8.0+cpu to override the default

backend/find_pytorch.py (5)

118-120: Update the comment to match the SpecifierSet (12.x, not just 12.2).

The code accepts any CUDA in [12,13), but the comment says “CUDA 12.2”. Recommend broadening the comment for accuracy.

-        if cuda_version == "" or cuda_version in SpecifierSet(">=12,<13"):
-            # CUDA 12.2, cudnn 9
+        if cuda_version == "" or cuda_version in SpecifierSet(">=12,<13"):
+            # CUDA 12.x (cudnn 9)
             pt_version = "2.8.0"

---

109-128: Respect explicit pt_version/env override; only default when unset.

Currently, the CIBUILDWHEEL branch unconditionally sets pt_version even if the caller passed an explicit pt_version or set PYTORCH_VERSION. Suggest reading the env override first and only defaulting when still empty. This preserves caller intent and reduces surprise.

     if pt_version is None:
         return {"torch": []}
-    if (
-        os.environ.get("CIBUILDWHEEL", "0") == "1"
-        and platform.system() == "Linux"
-        and platform.machine() == "x86_64"
-    ):
-        cuda_version = os.environ.get("CUDA_VERSION", "12.2")
-        if cuda_version == "" or cuda_version in SpecifierSet(">=12,<13"):
-            # CUDA 12.2, cudnn 9
-            pt_version = "2.8.0"
-        elif cuda_version in SpecifierSet(">=11,<12"):
-            # CUDA 11.8, cudnn 8
-            pt_version = "2.3.1"
-        else:
-            raise RuntimeError("Unsupported CUDA version") from None
     if pt_version == "":
         pt_version = os.environ.get("PYTORCH_VERSION", "")
+    if (
+        os.environ.get("CIBUILDWHEEL", "0") == "1"
+        and platform.system() == "Linux"
+        and platform.machine() == "x86_64"
+        and (pt_version == "" or pt_version.lower() == "auto")
+    ):
+        cuda_version = os.environ.get("CUDA_VERSION", "12.2")
+        if cuda_version == "" or cuda_version in SpecifierSet(">=12,<13"):
+            # CUDA 12.x (cudnn 9)
+            pt_version = "2.8.0"
+        elif cuda_version in SpecifierSet(">=11,<12"):
+            # CUDA 11.8, cudnn 8
+            pt_version = "2.3.1"
+        else:
+            raise RuntimeError("Unsupported CUDA version") from None

---

48-54: Fix docstring: “TensorFlow requirement” → “PyTorch requirement”.

Minor copy/paste oversight.

     Returns
     -------
     str, optional
         PyTorch library path if found.
     list of str
-        TensorFlow requirement if not found. Empty if found.
+        PyTorch requirement if not found. Empty if found.

---

149-151: Fix docstring: “TF” → “PyTorch”.

Clarify the function description to match the implementation.

-def get_pt_version(pt_path: Optional[Union[str, Path]]) -> str:
-    """Get TF version from a TF Python library path.
+def get_pt_version(pt_path: Optional[Union[str, Path]]) -> str:
+    """Get PyTorch version from a PyTorch Python library path.

---

164-168: Avoid executing arbitrary module code when reading version; add safe fallback.

Importing and exec’ing version.py executes code from an arbitrary path. Add a safe fallback that parses version without execution.

-    spec = importlib.util.spec_from_file_location("torch.version", version_file)
-    module = importlib.util.module_from_spec(spec)
-    spec.loader.exec_module(module)
-    return module.__version__
+    try:
+        spec = importlib.util.spec_from_file_location("torch.version", version_file)
+        module = importlib.util.module_from_spec(spec)  # type: ignore[arg-type]
+        if spec is None or spec.loader is None:
+            raise ImportError("Failed to load spec")
+        spec.loader.exec_module(module)  # type: ignore[union-attr]
+        return module.__version__  # type: ignore[attr-defined]
+    except Exception:
+        import re
+        text = version_file.read_text(encoding="utf-8")
+        m = re.search(r'__version__\s*=\s*[\'"]([^\'"]+)[\'"]', text)
+        return m.group(1) if m else ""

Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro