diff --git a/flake.lock b/flake.lock
index b40f51b3..14011768 100644
--- a/flake.lock
+++ b/flake.lock
@@ -944,11 +944,11 @@
         "nixpkgs": "nixpkgs_6"
       },
       "locked": {
-        "lastModified": 1724270760,
-        "narHash": "sha256-KX566x0+3HZcB20HPdvdwyMm7ZJg21M+iqVrs/HCimA=",
+        "lastModified": 1724784743,
+        "narHash": "sha256-NdEoWeNwR/ZstYnHaiQWIYZvr7VsrAh7g3+ZHUPrxuI=",
         "owner": "danieldk",
         "repo": "tgi-nix",
-        "rev": "12cbaa76ff258351741d3b5afb7161f617fe7b4c",
+        "rev": "c9580c3e39a855246bb87b584bbea1885b44f524",
         "type": "github"
       },
       "original": {
diff --git a/flake.nix b/flake.nix
index 83feb26a..0739c90a 100644
--- a/flake.nix
+++ b/flake.nix
@@ -31,15 +31,12 @@
           src = ./.;
           additionalCargoNixArgs = [ "--all-features" ];
         };
-        config = {
-          allowUnfree = true;
-          cudaSupport = true;
-        };
         pkgs = import nixpkgs {
-          inherit config system;
+          inherit system;
+          inherit (tgi-nix.lib) config;
           overlays = [
             rust-overlay.overlays.default
-            tgi-nix.overlay
+            tgi-nix.overlays.default
           ];
         };
         crateOverrides = import ./nix/crate-overrides.nix { inherit pkgs nix-filter; };
diff --git a/nix/server.nix b/nix/server.nix
index 4e0fdaa1..6ee088e0 100644
--- a/nix/server.nix
+++ b/nix/server.nix
@@ -28,6 +28,7 @@
   peft,
   safetensors,
   tokenizers,
+  torch,
   sentencepiece,
   transformers,
   typer,