diff --git a/src/main/scala/IF.scala b/src/main/scala/IF.scala
index b4fd54a..a08d526 100644
--- a/src/main/scala/IF.scala
+++ b/src/main/scala/IF.scala
@@ -46,9 +46,9 @@ class InstructionFetch extends MultiIOModule {
     * 
     * You should expand on or rewrite the code below.
     */
-  io.PC := PC
-  IMEM.io.instructionAddress := PC
-  PC := Mux(io.branch, io.branchAddress, Mux(io.stall, PC, PC + 4.U))
+  io.PC := Mux(io.branch, io.branchAddress, PC)
+  IMEM.io.instructionAddress := io.PC
+  PC := Mux(io.branch, io.branchAddress + 4.U, Mux(io.stall, PC, PC + 4.U))
 
   val instruction = Wire(new Instruction)
   instruction := IMEM.io.instruction.asTypeOf(new Instruction)
diff --git a/src/main/scala/IFBarrier.scala b/src/main/scala/IFBarrier.scala
index 4b5d103..789f4f6 100644
--- a/src/main/scala/IFBarrier.scala
+++ b/src/main/scala/IFBarrier.scala
@@ -20,22 +20,12 @@ class IFBarrier extends MultiIOModule {
 
   val instruction = Reg(new Instruction)
   val replay = RegInit(Bool(), false.B)
-  val flushRemaining = RegInit(UInt(2.W), 0.U)
-  flushRemaining := Mux(
-    io.flush,
-    1.U,
-    Mux(
-      flushRemaining === 0.U,
-      0.U,
-      flushRemaining - 1.U
-    )
-  )
 
   replay := io.stall
   instruction := io.instructionIn
 
   io.instructionOut := Mux(
-    io.stall || io.flush || flushRemaining > 0.U,
+    io.stall || io.flush,
     Instruction.NOP,
     Mux(
       replay,