feat(installer): use torch extra index on all cuda install pathways

invoke-ai · Oct 10, 2024 · eda9793 · eda9793
1 parent b6c7949
commit eda9793
Showing 1 changed file with 3 additions and 1 deletion.
diff --git a/installer/lib/installer.py b/installer/lib/installer.py
@@ -410,16 +410,18 @@ def get_torch_source() -> Tuple[str | None, str | None]:
             url = "https://download.pytorch.org/whl/rocm5.6"
         elif device == GpuType.CPU:
             url = "https://download.pytorch.org/whl/cpu"
-        # CUDA uses the default PyPi index
         elif device == GpuType.CUDA:
+            url = "https://download.pytorch.org/whl/cu124"
             optional_modules = "[onnx-cuda]"
         elif device == GpuType.CUDA_WITH_XFORMERS:
+            url = "https://download.pytorch.org/whl/cu124"
             optional_modules = "[xformers,onnx-cuda]"
     elif OS == "Windows":
         if device == GpuType.CUDA:
             url = "https://download.pytorch.org/whl/cu124"
             optional_modules = "[onnx-cuda]"
         elif device == GpuType.CUDA_WITH_XFORMERS:
+            url = "https://download.pytorch.org/whl/cu124"
             optional_modules = "[xformers,onnx-cuda]"
         elif device.value == "cpu":
             # CPU  uses the default PyPi index, no optional modules