diff --git a/compiler/src/iree/compiler/Dialect/Util/Transforms/OptimizeIntArithmetic.cpp b/compiler/src/iree/compiler/Dialect/Util/Transforms/OptimizeIntArithmetic.cpp
index d4b3a14b43f5..94b2ee237398 100644
--- a/compiler/src/iree/compiler/Dialect/Util/Transforms/OptimizeIntArithmetic.cpp
+++ b/compiler/src/iree/compiler/Dialect/Util/Transforms/OptimizeIntArithmetic.cpp
@@ -292,6 +292,13 @@ class DataFlowListener : public RewriterBase::Listener {
       s.eraseState(res);
   }
 
+  void notifyOperationModified(Operation *op) override {
+    s.eraseState(s.getProgramPointAfter(op));
+    for (Value res : op->getResults()) {
+      s.eraseState(res);
+    }
+  }
+
   DataFlowSolver &s;
 };