diff --git a/src/main/scala/components/hazard.scala b/src/main/scala/components/hazard.scala
index da0be412..ea074f5e 100644
--- a/src/main/scala/components/hazard.scala
+++ b/src/main/scala/components/hazard.scala
@@ -12,6 +12,9 @@ import chisel3._
  * Input:  idex_memread, true if the instruction in the ID/EX register is going to read from memory
  * Input:  idex_rd, the register number of the destination register for the instruction in the ID/EX register
  * Input:  exmem_taken, if true, then we are using the nextpc in the EX/MEM register, *not* pc+4.
+ * Input:  imem_good, high if instruction memory is idle and ready.
+  *        If either imem_good or dmem_good are false then CPU will stall
+ * Input:  dmem_good, high if data memory is idle and ready.
  *
  * Output: pcwrite, the value to write to the pc. If 0, pc+4, if 1 the next_pc from the memory stage,
  *         if 2, then the last pc value (2 stalls the pipeline)
@@ -31,6 +34,9 @@ class HazardUnit extends Module {
     val idex_rd      = Input(UInt(5.W))
     val exmem_taken  = Input(Bool())
 
+    val imem_good    = Input(Bool())
+    val dmem_good    = Input(Bool())
+
     val pcwrite      = Output(UInt(2.W))
     val ifid_bubble  = Output(Bool())
     val idex_bubble  = Output(Bool())
@@ -60,4 +66,11 @@ class HazardUnit extends Module {
     io.idex_bubble  := true.B
     io.exmem_bubble := true.B
   }
+
+  // mem stall
+  when (~(io.imem_good && io.dmem_good)) {
+    io.pcwrite     := 2.U // freeze PC
+    io.idex_bubble  := true.B
+    //io.exmem_bubble := true.B
+  }
 }
diff --git a/src/main/scala/pipelined/cpu.scala b/src/main/scala/pipelined/cpu.scala
index d59b5bbf..eb7f77ce 100644
--- a/src/main/scala/pipelined/cpu.scala
+++ b/src/main/scala/pipelined/cpu.scala
@@ -113,9 +113,6 @@ class PipelinedCPU(implicit val conf: CPUConfig) extends BaseCPU {
   // For wb back to other stages
   val write_data = Wire(UInt())
 
-  // Stall pipeline if neither instruction nor data memory are ready
-  val memStall = ~(io.imem.good && io.dmem.good)
-
   /////////////////////////////////////////////////////////////////////////////
   // FETCH STAGE
   /////////////////////////////////////////////////////////////////////////////
@@ -123,10 +120,13 @@ class PipelinedCPU(implicit val conf: CPUConfig) extends BaseCPU {
   // Note: This comes from the memory stage!
   // Only update the pc if the pcwrite flag is enabled
   if (conf.debug) { printf(p"PC: $pc\n") }
-  pc := MuxCase(0.U, Array(
+
+  val new_pc = MuxCase(0.U, Array(
                 (hazard.io.pcwrite === 0.U) -> pcPlusFour.io.result,
                 (hazard.io.pcwrite === 1.U) -> next_pc,
-                (hazard.io.pcwrite === 2.U) -> pc))
+                (hazard.io.pcwrite === 2.U || (io.imem.valid || io.dmem.valid)) -> pc))
+
+  pc := new_pc
 
   // Send the PC to the instruction memory port to get the instruction
   io.imem.address := pc
@@ -137,11 +137,10 @@ class PipelinedCPU(implicit val conf: CPUConfig) extends BaseCPU {
 
   // Fill the IF/ID register if we are not bubbling IF/ID
   // otherwise, leave the IF/ID register *unchanged*
-  when (~hazard.io.ifid_bubble && ~memStall) {
+  when (~hazard.io.ifid_bubble) {
     if_id.instruction := io.imem.instruction
     if_id.pc          := pc
     if_id.pcplusfour  := pcPlusFour.io.result
-
     io.imem.valid := true.B
   } .otherwise {
     io.imem.valid := false.B
@@ -152,6 +151,7 @@ class PipelinedCPU(implicit val conf: CPUConfig) extends BaseCPU {
     if_id.instruction := 0.U
     if_id.pc          := 0.U
     if_id.pcplusfour  := 0.U
+    io.imem.valid := false.B
   }
 
   if (conf.debug) { printf(p"IF/ID: $if_id\n") }
@@ -206,7 +206,7 @@ class PipelinedCPU(implicit val conf: CPUConfig) extends BaseCPU {
   id_ex.wbcontrol.toreg    := control.io.toreg
   id_ex.wbcontrol.regwrite := control.io.regwrite
 
-  when (hazard.io.idex_bubble || memStall) {
+  when (hazard.io.idex_bubble) {
     // Set the id_ex control to 0 to indicate a bubble
     id_ex.excontrol := 0.U.asTypeOf(new EXControl)
     id_ex.mcontrol  := 0.U.asTypeOf(new MControl)
@@ -298,7 +298,7 @@ class PipelinedCPU(implicit val conf: CPUConfig) extends BaseCPU {
   }
 
   // Check for bubble EX/MEM register
-  when (hazard.io.exmem_bubble || memStall) {
+  when (hazard.io.exmem_bubble) {
     // Set the ex_mem control to 0 to indicate a bubble
     ex_mem.mcontrol  := 0.U.asTypeOf(new MControl)
     ex_mem.wbcontrol := 0.U.asTypeOf(new WBControl)
@@ -328,17 +328,26 @@ class PipelinedCPU(implicit val conf: CPUConfig) extends BaseCPU {
 
   // Send input signals to the hazard detection unit
   hazard.io.exmem_taken := ex_mem.taken
+  // Send memory good signals to the hazard detection unit
+  hazard.io.imem_good := io.imem.good
+  hazard.io.dmem_good := io.dmem.good
+
 
   // Send input signals to the forwarding unit
   forwarding.io.exmemrd := ex_mem.writereg
   forwarding.io.exmemrw := ex_mem.wbcontrol.regwrite
 
-  // Wire the MEM/WB register
-  mem_wb.writereg   := ex_mem.writereg
-  mem_wb.aluresult  := ex_mem.aluresult
-  mem_wb.pcplusfour := ex_mem.pcplusfour
-  mem_wb.readdata   := io.dmem.readdata
-  mem_wb.wbcontrol  := ex_mem.wbcontrol
+  // Check for bubble in MEM/WB register
+  // This is a deviation from the CPU model described in Patterson & Hennessy.
+  // Primarily, the MEM/WB register is frozen when there is a memory stall
+  when (io.imem.good && io.dmem.good) {
+    // Wire the MEM/WB register
+    mem_wb.writereg := ex_mem.writereg
+    mem_wb.aluresult := ex_mem.aluresult
+    mem_wb.pcplusfour := ex_mem.pcplusfour
+    mem_wb.readdata := io.dmem.readdata
+    mem_wb.wbcontrol := ex_mem.wbcontrol
+  }
 
   if (conf.debug) { printf(p"MEM/WB: $mem_wb\n") }
 
@@ -352,6 +361,7 @@ class PipelinedCPU(implicit val conf: CPUConfig) extends BaseCPU {
                        (mem_wb.wbcontrol.toreg === 1.U) -> mem_wb.readdata,
                        (mem_wb.wbcontrol.toreg === 2.U) -> mem_wb.pcplusfour))
 
+
   // Write the data to the register file
   registers.io.writedata := write_data
   registers.io.writereg  := mem_wb.writereg
@@ -361,5 +371,7 @@ class PipelinedCPU(implicit val conf: CPUConfig) extends BaseCPU {
   forwarding.io.memwbrd := mem_wb.writereg
   forwarding.io.memwbrw := mem_wb.wbcontrol.regwrite
 
+  printf(p">>> Cycle #${cycleCount}\nPC:$pc (${if_id.pc})($new_pc)\n${id_ex.wbcontrol.toreg} => ${ex_mem.wbcontrol.toreg} => ${mem_wb.wbcontrol.toreg} (readdata = ${mem_wb.readdata})\n${!(io.imem.good && io.dmem.good)}\n\n")
+
   if (conf.debug) { printf("---------------------------------------------\n") }
 }