diff --git a/lib/mk-nixos.nix b/lib/mk-nixos.nix
index 5d1624c..90d1d03 100644
--- a/lib/mk-nixos.nix
+++ b/lib/mk-nixos.nix
@@ -11,6 +11,8 @@ let
     nix.registry.nixpkgs.flake = nixpkgs;
   };
   machineRoot = ../system/machines + "/${name}";
+  specificConfig = machineRoot + /configuration.nix;
+  diskoConfig = import (machineRoot + /disks.nix) { };
 in
 nixpkgs.lib.nixosSystem {
   inherit system specialArgs;
@@ -20,7 +22,7 @@ nixpkgs.lib.nixosSystem {
     disko.nixosModules.disko
     ./cachix.nix
     ../system/configuration.nix
-    (machineRoot + /configuration.nix)
-    (machineRoot + /disks.nix)
+    specificConfig
+    diskoConfig
   ];
 }
diff --git a/system/machines/vm/disks.nix b/system/machines/vm/disks.nix
index c639626..2eb913a 100644
--- a/system/machines/vm/disks.nix
+++ b/system/machines/vm/disks.nix
@@ -1,15 +1,19 @@
+{ disks ? [ "/dev/sda" ], ... }:
+let
+  mainDisk = builtins.elemAt disks 0;
+in
 {
   disko.devices = {
-    disk.sda = {
-      device = "/dev/sda";
+    disk.main = {
+      device = mainDisk;
       type = "disk";
       content = {
         type = "gpt";
         partitions = {
           ESP = {
             name = "ESP";
-            device = "/dev/disk/by-label/boot";
-            end = "512M";
+            type = "EF00";
+            size = "512M";
             content = {
               type = "filesystem";
               format = "vfat";
@@ -18,9 +22,7 @@
           };
           root = {
             name = "root";
-            device = "/dev/disk/by-label/nixos";
-            start = "512M";
-            end = "-10G";
+            end = "-2G";
             content = {
               type = "filesystem";
               format = "ext4";
@@ -28,8 +30,6 @@
             };
           };
           swap = {
-            name = "swap";
-            device = "/dev/disk/by-label/swap";
             size = "100%";
             content = {
               type = "swap";