diff --git a/.devops/nix/package.nix b/.devops/nix/package.nix
index cfffac257..5844b886a 100644
--- a/.devops/nix/package.nix
+++ b/.devops/nix/package.nix
@@ -1,5 +1,6 @@
 {
   lib,
+  autoAddDriverRunpath,
   glibc,
   config,
   stdenv,
diff --git a/flake.nix b/flake.nix
index c69637d11..a5617d37a 100644
--- a/flake.nix
+++ b/flake.nix
@@ -162,7 +162,7 @@
                 cuda = config.legacyPackages.llamaPackagesCuda.llama-cpp;
 
                 mpi-cpu = config.packages.default.override { useMpi = true; };
-                mpi-cuda = config.packages.default.override { useMpi = true; };
+                mpi-cuda = config.packages.cuda.override { useMpi = true; };
               }
               // lib.optionalAttrs (system == "x86_64-linux") {
                 rocm = config.legacyPackages.llamaPackagesRocm.llama-cpp;