diff --git a/module.nix b/module.nix
index 47dca00..4a5e070 100644
--- a/module.nix
+++ b/module.nix
@@ -192,41 +192,36 @@ in
     visible = "shallow";
   };
 
-  config = lib.mkMerge [
-    (lib.mkIf (cfg.devices.disk != { }) {
-      system.build = (cfg.devices._scripts { inherit pkgs; checked = cfg.checkScripts; }) // {
+  config = {
+    _module.args.diskoLib = import ./lib {
+      inherit lib;
+      rootMountPoint = config.disko.rootMountPoint;
+      makeTest = import (pkgs.path + "/nixos/tests/make-test-python.nix");
+      eval-config = import (pkgs.path + "/nixos/lib/eval-config.nix");
+    };
 
-        # we keep these old outputs for compatibility
-        disko = builtins.trace "the .disko output is deprecated, please use .diskoScript instead" (cfg.devices._scripts { inherit pkgs; }).diskoScript;
-        diskoNoDeps = builtins.trace "the .diskoNoDeps output is deprecated, please use .diskoScriptNoDeps instead" (cfg.devices._scripts { inherit pkgs; }).diskoScriptNoDeps;
+    system.build = (cfg.devices._scripts { inherit pkgs; checked = cfg.checkScripts; }) // {
+      # we keep these old outputs for compatibility
+      disko = builtins.trace "the .disko output is deprecated, please use .diskoScript instead" (cfg.devices._scripts { inherit pkgs; }).diskoScript;
+      diskoNoDeps = builtins.trace "the .diskoNoDeps output is deprecated, please use .diskoScriptNoDeps instead" (cfg.devices._scripts { inherit pkgs; }).diskoScriptNoDeps;
 
-        installTest = diskoLib.testLib.makeDiskoTest {
-          inherit extendModules pkgs;
-          name = "${config.networking.hostName}-disko";
-          disko-config = builtins.removeAttrs config [ "_module" ];
-          testMode = "direct";
-          efi = cfg.tests.efi;
-          extraSystemConfig = cfg.tests.extraConfig;
-          extraTestScript = cfg.tests.extraChecks;
-        };
-
-        vmWithDisko = lib.mkDefault config.virtualisation.vmVariantWithDisko.system.build.vmWithDisko;
+      installTest = diskoLib.testLib.makeDiskoTest {
+        inherit extendModules pkgs;
+        name = "${config.networking.hostName}-disko";
+        disko-config = builtins.removeAttrs config [ "_module" ];
+        testMode = "direct";
+        efi = cfg.tests.efi;
+        extraSystemConfig = cfg.tests.extraConfig;
+        extraTestScript = cfg.tests.extraChecks;
       };
 
+      vmWithDisko = lib.mkDefault config.virtualisation.vmVariantWithDisko.system.build.vmWithDisko;
+    };
 
-      # we need to specify the keys here, so we don't get an infinite recursion error
-      # Remember to add config keys here if they are added to types
-      fileSystems = lib.mkIf cfg.enableConfig cfg.devices._config.fileSystems or { };
-      boot = lib.mkIf cfg.enableConfig cfg.devices._config.boot or { };
-      swapDevices = lib.mkIf cfg.enableConfig cfg.devices._config.swapDevices or [ ];
-    })
-    {
-      _module.args.diskoLib = import ./lib {
-        inherit lib;
-        rootMountPoint = config.disko.rootMountPoint;
-        makeTest = import (pkgs.path + "/nixos/tests/make-test-python.nix");
-        eval-config = import (pkgs.path + "/nixos/lib/eval-config.nix");
-      };
-    }
-  ];
+    # we need to specify the keys here, so we don't get an infinite recursion error
+    # Remember to add config keys here if they are added to types
+    fileSystems = lib.mkIf cfg.enableConfig cfg.devices._config.fileSystems or { };
+    boot = lib.mkIf cfg.enableConfig cfg.devices._config.boot or { };
+    swapDevices = lib.mkIf cfg.enableConfig cfg.devices._config.swapDevices or [ ];
+  };
 }