diff --git a/src/tir/transforms/bf16_legalize.cc b/src/tir/transforms/bf16_legalize.cc
index 07f4775ded50..445ac1cf60cd 100644
--- a/src/tir/transforms/bf16_legalize.cc
+++ b/src/tir/transforms/bf16_legalize.cc
@@ -330,7 +330,7 @@ class BF16LowerRewriter : StmtExprMutator {
       }
     }
     if (buffer_remap.size() != 0) {
-      op->buffer_map.assign(changes.begin(), changes.end());
+      op->buffer_map = Map<Var, Buffer>(changes.begin(), changes.end());
     }
   }
 };