diff --git a/pkgs/development/python-modules/onnxscript/default.nix b/pkgs/development/python-modules/onnxscript/default.nix
index 0cce49403530..a8e6baec1378 100644
--- a/pkgs/development/python-modules/onnxscript/default.nix
+++ b/pkgs/development/python-modules/onnxscript/default.nix
@@ -15,7 +15,8 @@
   onnx-ir,
   packaging,
   typing-extensions,
-  pynvml,
+  # cuda-only:
+  nvidia-ml-py,
 
   # tests
   onnxruntime,
@@ -58,7 +59,7 @@ buildPythonPackage (finalAttrs: {
     typing-extensions
   ]
   ++ lib.optionals cudaSupport [
-    pynvml
+    nvidia-ml-py
   ];
 
   pythonImportsCheck = [ "onnxscript" ];