[MLIR:TF] Fix invalid access on SmallDenseMap problem in TensorList decomposition.

PiperOrigin-RevId: 306553065
Change-Id: Iec565d64cf9f7b5c2277cda43edeb61a1a067242

tensorflow/compiler/mlir/tensorflow/transforms/tensor_list_ops_decomposition.cc

@@ -541,7 +541,8 @@ LogicalResult HandleTensorListSetItemOp(
   auto new_buffer = cutil::SetElement(index, buffer, set_item.item(), builder,
                                       set_item.getLoc());
   set_item.output_handle().replaceAllUsesWith(new_buffer);
-  (*buffer_to_size)[new_buffer] = it->getSecond();
+  auto size = it->getSecond();
+  (*buffer_to_size)[new_buffer] = size;
   set_item.erase();
   return success();
 }
@@ -665,12 +666,14 @@ LogicalResult DecomposeTensorListOpsInternal(
       auto it = buffer_to_size->find(addn.getOperand(0));
       if (it != buffer_to_size->end()) {
         addn.sum().setType(addn.getOperand(0).getType());
-        (*buffer_to_size)[addn.sum()] = it->getSecond();
+        auto size = it->getSecond();
+        (*buffer_to_size)[addn.sum()] = size;
       }
     } else if (auto zeros = llvm::dyn_cast<TF::ZerosLikeOp>(&op)) {
       if (buffer_to_size->count(zeros.x()) > 0) {
         zeros.y().setType(zeros.x().getType());
-        (*buffer_to_size)[zeros.y()] = (*buffer_to_size)[zeros.x()];
+        auto size = (*buffer_to_size)[zeros.x()];
+        (*buffer_to_size)[zeros.y()] = size;
       }
     } else if (auto while_op = llvm::dyn_cast<TF::WhileOp>(&op)) {
       if (failed(HandleWhileOp(while_op, module, buffer_to_size,