diff --git a/conda/environments/all_cuda-121_arch-x86_64.yaml b/conda/environments/all_cuda-121_arch-x86_64.yaml
index 77bd861..825aa8c 100644
--- a/conda/environments/all_cuda-121_arch-x86_64.yaml
+++ b/conda/environments/all_cuda-121_arch-x86_64.yaml
@@ -10,6 +10,7 @@ dependencies:
 - breathe
 - cmake>=3.26.4,!=3.30.0
 - cuda-cudart-dev
+- cuda-nvml-dev
 - cuda-nvtx-dev
 - cuda-profiler-api
 - cuda-version=12.1
diff --git a/conda/environments/all_cuda-124_arch-x86_64.yaml b/conda/environments/all_cuda-124_arch-x86_64.yaml
index 30a259e..69e24e6 100644
--- a/conda/environments/all_cuda-124_arch-x86_64.yaml
+++ b/conda/environments/all_cuda-124_arch-x86_64.yaml
@@ -10,6 +10,7 @@ dependencies:
 - breathe
 - cmake>=3.26.4,!=3.30.0
 - cuda-cudart-dev
+- cuda-nvml-dev
 - cuda-nvtx-dev
 - cuda-profiler-api
 - cuda-version=12.4
@@ -46,7 +47,7 @@ dependencies:
 - pytest-cov
 - pytest-forked
 - pytest-xdist
-- pytorch-cuda=12.1
+- pytorch-cuda=12.4
 - pytorch::pytorch>=2.3,<2.4.0a0
 - pytorch_geometric>=2.5,<2.6
 - raft-dask==24.12.*,>=0.0.0a0
diff --git a/dependencies.yaml b/dependencies.yaml
index 1787b05..8cb55d1 100644
--- a/dependencies.yaml
+++ b/dependencies.yaml
@@ -244,6 +244,7 @@ dependencies:
               cuda: "12.*"
             packages:
               - cuda-cudart-dev
+              - cuda-nvml-dev
               - cuda-nvtx-dev
               - cuda-profiler-api
               - libcublas-dev
@@ -413,12 +414,19 @@ dependencies:
           - {matrix: null, packages: [*pytorch_pip, *tensordict]}
       - output_types: [conda]
         matrices:
-          - matrix: {cuda: "12.*"}
+          - matrix: {cuda: "12.1"}
             packages:
               - pytorch-cuda=12.1
-          - matrix: {cuda: "11.*"}
+          - matrix: {cuda: "12.4"}
+            packages:
+              - pytorch-cuda=12.4
+          - matrix: {cuda: "11.8"}
             packages:
               - pytorch-cuda=11.8
+          # pytorch only supports certain CUDA versions... skip
+          # adding pytorch-cuda pinning if any other CUDA version is requested
+          - matrix:
+            packages:
 
   depends_on_dgl:
     specific: