python3Packages.torchcodec: respect torch.cudaCapabilities (#530330)

2026-06-10 14:21:53 +00:00
parent db0a1574b8 887d12e845
commit 07edef9b4d
1 changed files with 3 additions and 0 deletions
@@ -99,6 +99,9 @@ buildPythonPackage.override { inherit (torch) stdenv; } (finalAttrs: {

    ENABLE_CUDA = cudaSupport;
  }
+  // lib.optionalAttrs cudaSupport {
+    TORCH_CUDA_ARCH_LIST = "${lib.concatStringsSep ";" torch.cudaCapabilities}";
+  }
  // lib.optionalAttrs rocmSupport {
    ROCM_PATH = torch.rocmtoolkit_joined;
    ROCM_SOURCE_DIR = torch.rocmtoolkit_joined;