diff --git a/mlir/test/lib/Dialect/Transform/TestTransformDialectInterpreter.cpp b/mlir/test/lib/Dialect/Transform/TestTransformDialectInterpreter.cpp
index 1273414cd4dfc..bd2de68fb276d 100644
--- a/mlir/test/lib/Dialect/Transform/TestTransformDialectInterpreter.cpp
+++ b/mlir/test/lib/Dialect/Transform/TestTransformDialectInterpreter.cpp
@@ -36,13 +36,18 @@ struct TestTransformDialectEraseSchedulePass
   }
 
   void runOnOperation() override {
+    SmallVector<Operation *> opsToDelete;
     getOperation()->walk<WalkOrder::PreOrder>([&](Operation *nestedOp) {
       if (isa<transform::TransformOpInterface>(nestedOp)) {
-        nestedOp->erase();
+        opsToDelete.push_back(nestedOp);
         return WalkResult::skip();
       }
       return WalkResult::advance();
     });
+    for (Operation *op : llvm::reverse(opsToDelete)) {
+      // erase the operation
+      op->erase();
+    }
   }
 };
 } // namespace