diff --git a/mlir/lib/Dialect/MPI/IR/MPIOps.cpp b/mlir/lib/Dialect/MPI/IR/MPIOps.cpp
index e5e09e28998ba..3c93029f0d37f 100644
--- a/mlir/lib/Dialect/MPI/IR/MPIOps.cpp
+++ b/mlir/lib/Dialect/MPI/IR/MPIOps.cpp
@@ -52,7 +52,7 @@ struct FoldCast final : public mlir::OpRewritePattern<OpT> {
     if (!src.getType().hasStaticShape()) {
       return mlir::failure();
     }
-    op.getRefMutable().assign(src);
+    b.modifyOpInPlace(op, [&]() { op.getRefMutable().assign(src); });
     return mlir::success();
   }
 };