1 files changed, 1 insertions, 1 deletions
diff --git a/mlir/lib/Dialect/Arith/Transforms/IntRangeOptimizations.cpp b/mlir/lib/Dialect/Arith/Transforms/IntRangeOptimizations.cpp
index 8005f9103b23..521138c1f6f4 100644
--- a/mlir/lib/Dialect/Arith/Transforms/IntRangeOptimizations.cpp
+++ b/mlir/lib/Dialect/Arith/Transforms/IntRangeOptimizations.cpp
@@ -80,7 +80,7 @@ public:
 
 protected:
   void notifyOperationErased(Operation *op) override {
-    s.eraseState(op);
+    s.eraseState(s.getProgramPointAfter(op));
     for (Value res : op->getResults())
       s.eraseState(res);
   }