diff --git a/src/stages/partition/mod.rs b/src/stages/partition/mod.rs
index 477fa2c..3f06070 100644
--- a/src/stages/partition/mod.rs
+++ b/src/stages/partition/mod.rs
@@ -156,11 +156,17 @@ impl PartitionStage {
     }
 
     fn get_swap_size(&self) -> KiraSize {
+        if let Some(s_mode) = self.swap_mode
+            && s_mode == SwapMode::SwapCustom
+        {
+            return KiraSize::new_mb(self.custom_swap_size_mb as u64);
+        }
+
         let total_ram = self.system_ram_size + self.vram_of_all_gpus_size;
-        log::debug!("total_ram {}", total_ram);
+
         let sys_ram_gb = self.system_ram_size.gb();
         log::debug!("sys_ram_gb {}", sys_ram_gb);
-        let disk_swap_gb = if sys_ram_gb >= 32 {
+        let mut disk_swap_gb = if sys_ram_gb >= 32 {
             1u64
         } else if sys_ram_gb >= 24 {
             2u64
@@ -172,9 +178,11 @@ impl PartitionStage {
             sys_ram_gb * 2
         };
 
-        log::debug!("disk_swap_gb {}", disk_swap_gb);
+        if self.use_zram {
+            disk_swap_gb = disk_swap_gb.saturating_sub(sys_ram_gb).max(2);
+        }
         let disk_swap = KiraSize::new_gb(disk_swap_gb);
-        log::debug!("disk_swap {}", disk_swap);
+
         match self.swap_mode {
             Some(SwapMode::NoSwap) => KiraSize::new_b(0),
             Some(SwapMode::SwapHibernate) => disk_swap + total_ram,
@@ -183,6 +191,19 @@ impl PartitionStage {
         }
     }
 
+    fn gen_layout(&mut self) {
+        if let Some(dev) = self.device.clone() {
+            let swap_size: KiraSize = self.get_swap_size();
+            self.generated_disk_parts = Some(
+                PartLayout::gen_classic_layout(dev.clone(), KiraSize::new_mb(512), swap_size)
+                    .part_list
+                    .iter()
+                    .map(|part_info| part_to_ct_params(part_info, dev.size))
+                    .collect(),
+            );
+        }
+    }
+
     fn gen_result(&self) -> StageResult {
         StageResult::new("partition")
             .add_val_string(
@@ -236,22 +257,7 @@ impl PartitionStage {
 
             Message::SelectSwapMode(sm) => {
                 self.swap_mode = Some(sm);
-
-                if let Some(dev) = self.device.clone() {
-                    let swap_size: KiraSize = self.get_swap_size();
-                    self.generated_disk_parts = Some(
-                        PartLayout::gen_classic_layout(
-                            dev.clone(),
-                            KiraSize::new_mb(512),
-                            swap_size,
-                        )
-                        .part_list
-                        .iter()
-                        .map(|part_info| part_to_ct_params(part_info, dev.size))
-                        .collect(),
-                    );
-                }
-
+                self.gen_layout();
                 StageAction::None
             }
             Message::ToggleSecureBoot(t) => {
@@ -260,6 +266,7 @@ impl PartitionStage {
             }
             Message::ToggleZram(t) => {
                 self.use_zram = t;
+                self.gen_layout();
                 StageAction::None
             }
             Message::SwapSizeIncrement => {
@@ -267,6 +274,7 @@ impl PartitionStage {
                     .strict_mul(1024u32)
                     .saturating_add(1024u32);
                 self.custom_swap_size_str = format!("{}MB", self.custom_swap_size_mb);
+                self.gen_layout();
                 StageAction::None
             }
             Message::SwapSizeDecrement => {
@@ -274,6 +282,7 @@ impl PartitionStage {
                     .strict_mul(1024u32)
                     .saturating_sub(1024u32);
                 self.custom_swap_size_str = format!("{}MB", self.custom_swap_size_mb);
+                self.gen_layout();
                 StageAction::None
             }
             Message::SwapSizeFieldChange(f_v) => {
@@ -281,6 +290,7 @@ impl PartitionStage {
                 only_num.retain(char::is_numeric);
                 self.custom_swap_size_mb = u32::from_str_radix(&only_num, 10).unwrap_or(1024);
                 self.custom_swap_size_str = format!("{}MB", self.custom_swap_size_mb);
+                self.gen_layout();
                 StageAction::None
             }
             Message::Back => StageAction::Back,