python3Packages.torchcodec: properly target cuda capabilities

2026-06-10 12:12:01 +00:00
parent 013b417c09
commit 887d12e845
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;