diff --git a/hosts/nixos-desktop/disko.nix b/hosts/nixos-desktop/disko.nix
index d373c43..43d1242 100644
--- a/hosts/nixos-desktop/disko.nix
+++ b/hosts/nixos-desktop/disko.nix
@@ -1,4 +1,4 @@
-{ lib, ... }:
+{ lib, pkgs, ... }:
 
 {
   # needed for zfs pool
@@ -83,9 +83,27 @@
     };
   };
   # rollback to blank
-  boot.initrd.postDeviceCommands = lib.mkAfter ''
-    zfs rollback -r zroot/encrypted/root@blank && echo "rollback complete"
-  '';
+  boot.initrd.systemd.services.rollback = {
+    description = "Rollback ZFS datasets to a pristine state";
+    wantedBy = [
+      "initrd.target"
+    ];
+    after = [
+      "zfs-import-zroot.service"
+    ];
+    before = [
+      "sysroot.mount"
+    ];
+    path = with pkgs; [
+      zfs
+    ];
+    unitConfig.DefaultDependencies = "no";
+    serviceConfig.Type = "oneshot";
+    script = ''
+      zfs rollback -r zroot/encrypted/root@blank && echo "rollback complete"
+    '';
+  };
+  boot.initrd.systemd.enable = true;
   fileSystems."/persist".neededForBoot = true;
   # HACK: to fix issue of agenix running before impermanence
   age.identityPaths = [