diff --git a/flake.lock b/flake.lock
index 2702dd4..e66988b 100644
--- a/flake.lock
+++ b/flake.lock
@@ -757,6 +757,21 @@
         "type": "github"
       }
     },
+    "nixos-nixpkgs": {
+      "locked": {
+        "lastModified": 1724956629,
+        "narHash": "sha256-ZX4gCiRd9VnfDDetgu02VIYvPTBPvfzNK53OJ3aGab0=",
+        "owner": "nixos",
+        "repo": "nixpkgs",
+        "rev": "cc223bad7ef2ad0ac4c2db75136864233e40465b",
+        "type": "github"
+      },
+      "original": {
+        "owner": "nixos",
+        "repo": "nixpkgs",
+        "type": "github"
+      }
+    },
     "nixpkgs": {
       "locked": {
         "lastModified": 1714555275,
@@ -1042,6 +1057,7 @@
         "flake-compat": "flake-compat",
         "flake-utils": "flake-utils",
         "naersk": "naersk",
+        "nixos-nixpkgs": "nixos-nixpkgs",
         "nixpkgs": "nixpkgs_2",
         "serokell-nix": "serokell-nix"
       }
diff --git a/flake.nix b/flake.nix
index c0d9b5c..2d392d0 100644
--- a/flake.nix
+++ b/flake.nix
@@ -11,14 +11,15 @@
   inputs = {
     flake-compat.flake = false;
     naersk.url = "github:nix-community/naersk";
+    nix.url = "github:nixos/nix?ref=2.21.4";
   };
 
-  outputs = { self, nixpkgs, flake-utils, serokell-nix, naersk, ... }:
+  outputs = { self, nixpkgs, flake-utils, serokell-nix, naersk, ... }@inputs:
     flake-utils.lib.eachSystem [ "x86_64-linux" ] (system:
       let
         pkgs = nixpkgs.legacyPackages.${system}.extend serokell-nix.overlay;
         naersk' = pkgs.callPackage naersk {};
-        nix = pkgs.nix;
+        nix = inputs.nix.packages.${system}.nix;
 
         update-daemon = naersk'.buildPackage {
           src = builtins.path {
diff --git a/src/main.rs b/src/main.rs
index 8aed8ad..cb535a1 100644
--- a/src/main.rs
+++ b/src/main.rs
@@ -51,21 +51,16 @@ fn flake_update(
     lock: &Lock,
 ) -> Result<(), FlakeUpdateError> {
     let mut nix_flake_update = Command::new("nix");
-    nix_flake_update.arg("flake");
+    nix_flake_update.arg("flake").arg("update");
 
-    // If a list of inputs to update is not provided, update all inputs
-    if settings.inputs.is_empty() {
-        nix_flake_update.arg("update");
-    // Otherwise, update only the specified inputs
-    } else {
-        nix_flake_update.arg("lock");
+    // If a list of inputs to update is provided, update only the specified inputs
+    if !settings.inputs.is_empty() {
         for input in settings.inputs.iter() {
             // Abort flake update if input is missing from the flake.lock root nodes
             // and allow_missing_inputs is not set
             if !settings.allow_missing_inputs && lock.get_root_dep(input.clone()).is_none() {
                 return Err(FlakeUpdateError::MissingInput(input.clone()));
             };
-            nix_flake_update.arg("--update-input");
             nix_flake_update.arg(input);
         }
     };