diff --git a/boot/grub/first.grub.cfg b/boot/grub/first.grub.cfg
index 8fe76fe..d2ae8f4 100644
--- a/boot/grub/first.grub.cfg
+++ b/boot/grub/first.grub.cfg
@@ -19,8 +19,7 @@ if [ "${rugpi_boot_spare}" = "true" ]; then
     set rugpi_boot_spare=true
 fi
 
-# TODO: This should not be hardcoded to `hd0`, either use `search` or extract the correct
-# root drive from the existing value of `root` in some way.
-set root=(hd0,${rugpi_bootpart})
+regexp (.*),(.*) $root -s 1:rugpi_root_device
+set root=(${rugpi_root_device},${rugpi_bootpart})
 
 source (${root})/second.grub.cfg