From 8bef8364976c7b396f1fa99442306b15025c6d3c Mon Sep 17 00:00:00 2001
From: "A. Unique TensorFlower" <gardener@tensorflow.org>
Date: Wed, 31 Jul 2024 01:56:19 -0700
Subject: [PATCH] Automated Code Change

PiperOrigin-RevId: 657916008
---
 xla/service/transfer_manager.cc | 3 +++
 1 file changed, 3 insertions(+)

diff --git a/xla/service/transfer_manager.cc b/xla/service/transfer_manager.cc
index b5edfad998ad5..c601a884919d1 100644
--- a/xla/service/transfer_manager.cc
+++ b/xla/service/transfer_manager.cc
@@ -270,6 +270,8 @@ absl::Status TransferManager::WriteRootTupleIndexTable(
                device_memory.size());
 
   std::vector<se::DeviceMemoryBase> elements;
+  elements.reserve(
+      ShapeUtil::TupleElementCount(device_buffer.on_device_shape()));
   for (int64_t i = 0;
        i < ShapeUtil::TupleElementCount(device_buffer.on_device_shape()); ++i) {
     elements.push_back(device_buffer.buffer({i}));
@@ -290,6 +292,7 @@ absl::Status TransferManager::WriteRootTupleIndexTable(
                device_memory.size());
 
   std::vector<se::DeviceMemoryBase> elements;
+  elements.reserve(ShapeUtil::TupleElementCount(buffer_tree.shape()));
   for (int64_t i = 0; i < ShapeUtil::TupleElementCount(buffer_tree.shape());
        ++i) {
     elements.push_back(buffer_tree.element({i}).AsDeviceMemoryBase());