diff --git a/src/main/scala/vexriscv/plugin/MulDivIterativePlugin.scala b/src/main/scala/vexriscv/plugin/MulDivIterativePlugin.scala
index f366854..11af4ab 100644
--- a/src/main/scala/vexriscv/plugin/MulDivIterativePlugin.scala
+++ b/src/main/scala/vexriscv/plugin/MulDivIterativePlugin.scala
@@ -69,8 +69,9 @@ class MulDivIterativePlugin(genMul : Boolean = true,
     import pipeline.config._
     if(!genMul && !genDiv) return
 
-    memory plug new Area {
-      import memory._
+    val flushStage = if(memory != null) memory else execute
+    flushStage plug new Area {
+      import flushStage._
 
       //Shared ressources
       val rs1 = Reg(UInt(33 bits))