diff --git a/pkgs/development/python-modules/torch/source/default.nix b/pkgs/development/python-modules/torch/source/default.nix
index 3f006e3e4384..f8d9d63a873b 100644
--- a/pkgs/development/python-modules/torch/source/default.nix
+++ b/pkgs/development/python-modules/torch/source/default.nix
@@ -14,7 +14,7 @@
   autoAddDriverRunpath,
   effectiveMagma ?
     if cudaSupport then
-      magma-cuda-static
+      magma-cuda-static.override { inherit cudaPackages; }
     else if rocmSupport then
       magma-hip
     else
@@ -81,7 +81,7 @@
   #          (dependencies without cuda support).
   #          Instead we should rely on overlays and nixpkgsFun.
   # (@SomeoneSerge)
-  _tritonEffective ? if cudaSupport then triton-cuda else triton,
+  _tritonEffective ? if cudaSupport then triton-cuda.override { inherit cudaPackages; } else triton,
   triton-cuda,
 
   # Disable MKLDNN on aarch64-darwin, it negatively impacts performance,
@@ -263,6 +263,9 @@ let
     "Magma cudaPackages does not match cudaPackages" =
       cudaSupport
       && (effectiveMagma.cudaPackages.cudaMajorMinorVersion != cudaPackages.cudaMajorMinorVersion);
+    "Triton cudaPackages does not match cudaPackages" =
+      cudaSupport
+      && (_tritonEffective.cudaPackages.cudaMajorMinorVersion != cudaPackages.cudaMajorMinorVersion);
   };
 
   unroll-src = writeShellScript "unroll-src" ''
diff --git a/pkgs/development/python-modules/torchaudio/default.nix b/pkgs/development/python-modules/torchaudio/default.nix
index 974bf23e7ec1..0db400402f72 100644
--- a/pkgs/development/python-modules/torchaudio/default.nix
+++ b/pkgs/development/python-modules/torchaudio/default.nix
@@ -28,12 +28,12 @@
   torchaudio,
 
   cudaSupport ? torch.cudaSupport,
-  cudaPackages,
   rocmSupport ? torch.rocmSupport,
 }:
 
 let
   inherit (stdenv) hostPlatform;
+  inherit (torch) cudaCapabilities cudaPackages;
 in
 buildPythonPackage.override { inherit (torch) stdenv; } (finalAttrs: {
   pname = "torchaudio";
@@ -53,7 +53,7 @@ buildPythonPackage.override { inherit (torch) stdenv; } (finalAttrs: {
 
     # CUDA
     USE_CUDA = cudaSupport;
-    TORCH_CUDA_ARCH_LIST = "${lib.concatStringsSep ";" torch.cudaCapabilities}";
+    TORCH_CUDA_ARCH_LIST = "${lib.concatStringsSep ";" cudaCapabilities}";
     TORCHAUDIO_TEST_ALLOW_SKIP_IF_NO_CUCTC_DECODER = 1;
     TORCHAUDIO_TEST_ALLOW_SKIP_IF_NO_CUDA = 1;
     TORCHAUDIO_TEST_ALLOW_SKIP_IF_NO_MULTIGPU_CUDA = 1;
diff --git a/pkgs/development/python-modules/torchcodec/default.nix b/pkgs/development/python-modules/torchcodec/default.nix
index db1d474a6be0..8523ebbe076c 100644
--- a/pkgs/development/python-modules/torchcodec/default.nix
+++ b/pkgs/development/python-modules/torchcodec/default.nix
@@ -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;
diff --git a/pkgs/development/python-modules/triton/default.nix b/pkgs/development/python-modules/triton/default.nix
index 72e8df58f076..92c3d8a25d38 100644
--- a/pkgs/development/python-modules/triton/default.nix
+++ b/pkgs/development/python-modules/triton/default.nix
@@ -218,6 +218,7 @@ buildPythonPackage.override { stdenv = effectiveStdenv; } (finalAttrs: {
   ];
 
   passthru = {
+    inherit cudaPackages;
     gpuCheck = effectiveStdenv.mkDerivation {
       pname = "triton-pytest";
       inherit (triton) version src;