diff --git a/src/main.rs b/src/main.rs
index 4cbaff9..019e4e9 100644
--- a/src/main.rs
+++ b/src/main.rs
@@ -50,8 +50,15 @@ fn format_drive(drive_name: &str) -> rsfdisk::Result<()> {
 }
 
 fn format_partitions(drive_name: &str) {
-    let efi_path = format!("{}p1", drive_name);
-    let root_path = format!("{}p2", drive_name);
+    // Check if drive name ends with a digit (e.g., "nvme0n1")
+    let suffix = if drive_name.chars().last().map(|c| c.is_ascii_digit()).unwrap_or(false) {
+        "p"
+    } else {
+        ""
+    };
+
+    let efi_path = format!("{}{}1", drive_name, suffix);
+    let root_path = format!("{}{}2", drive_name, suffix);
 
     let _efi_partition = Command::new("mkfs.fat")
         .arg("-F32")
@@ -68,8 +75,15 @@ fn format_partitions(drive_name: &str) {
 }
 
 fn mount_partitions(drive_name: &str) {
-    let efi_path = format!("{}p1", drive_name);
-    let root_path = format!("{}p2", drive_name);
+    // Check if drive name ends with a digit (e.g., "nvme0n1")
+    let suffix = if drive_name.chars().last().map(|c| c.is_ascii_digit()).unwrap_or(false) {
+        "p"
+    } else {
+        ""
+    };
+
+    let efi_path = format!("{}{}1", drive_name, suffix);
+    let root_path = format!("{}{}2", drive_name, suffix);
 
     let root_source = Some(Path::new(&root_path));
     let root_target = Path::new("/mnt");