nathan/nixpkgs
1
0
Fork 0
mirror of https://github.com/NixOS/nixpkgs.git synced 2026-07-21 08:01:31 +00:00
Code Issues Packages Projects Releases Wiki Activity

python314Packages.torchcodec: use cudaPackages from torch

Browse Source
This commit is contained in:
Sandro Jäckel
2026-07-12 16:57:00 +02:00
parent 3e54bbc91d
commit 41f3bd769c
1 changed files with 4 additions and 2 deletions
Download Patch File Download Diff File Expand all files Collapse all files

6
pkgs/development/python-modules/torchcodec/default.nix
View File

@@ -20,10 +20,12 @@
torchvision,
cudaSupport ? torch.cudaSupport,
cudaPackages,
rocmSupport ? torch.rocmSupport,
}:
let
inherit (torch) cudaCapabilities cudaPackages;
in
buildPythonPackage.override { inherit (torch) stdenv; } (finalAttrs: {
pname = "torchcodec";
version = "0.14.0";
@@ -100,7 +102,7 @@ buildPythonPackage.override { inherit (torch) stdenv; } (finalAttrs: {
ENABLE_CUDA = cudaSupport;
}
// lib.optionalAttrs cudaSupport {
TORCH_CUDA_ARCH_LIST = "${lib.concatStringsSep ";" torch.cudaCapabilities}";
TORCH_CUDA_ARCH_LIST = "${lib.concatStringsSep ";" cudaCapabilities}";
}
// lib.optionalAttrs rocmSupport {
ROCM_PATH = torch.rocmtoolkit_joined;