diff --git a/pkgs/development/python-modules/xformers/default.nix b/pkgs/development/python-modules/xformers/default.nix
index 1ad3c89c05c3..323159a80690 100644
--- a/pkgs/development/python-modules/xformers/default.nix
+++ b/pkgs/development/python-modules/xformers/default.nix
@@ -13,7 +13,7 @@
 
   # dependencies
   numpy,
-  pynvml,
+  nvidia-ml-py,
 
   # tests
   einops,
@@ -102,7 +102,7 @@ buildPythonPackage.override { stdenv = effectiveStdenv; } (finalAttrs: {
     torch
   ]
   ++ lib.optionals cudaSupport [
-    pynvml
+    nvidia-ml-py
   ];
 
   pythonImportsCheck = [