diff --git a/src/main/scala/MEM.scala b/src/main/scala/MEM.scala
index a136fa4..dd641c1 100644
--- a/src/main/scala/MEM.scala
+++ b/src/main/scala/MEM.scala
@@ -48,12 +48,6 @@ class MemoryFetch() extends MultiIOModule {
   DMEM.io.dataIn      := io.writeData
   DMEM.io.writeEnable := io.writeMem
 
-  when(io.writeMem) {
-    DMEM.io.dataAddress := ALUResult
-  }.otherwise(DMEM.io.dataAddress := io.ALUResult)
-
-
-  when(io.readMem) {
-    io.dataOut := DMEM.io.dataOut
-  }.otherwise(io.dataOut := ALUResult)
+  DMEM.io.dataAddress := Mux(io.writeMem, ALUResult, io.ALUResult)
+  io.dataOut := Mux(io.readMem, DMEM.io.dataOut, ALUResult)
 }