diff --git a/tensorflow/stream_executor/tpu/tpu_transfer_manager.cc b/tensorflow/stream_executor/tpu/tpu_transfer_manager.cc index e6b3460ac92..b49a0b56259 100644 --- a/tensorflow/stream_executor/tpu/tpu_transfer_manager.cc +++ b/tensorflow/stream_executor/tpu/tpu_transfer_manager.cc @@ -93,6 +93,8 @@ Status TpuTransferManager::TransferLiteralToInfeed( tpu::ExecutorApiFn()->TpuTransferManager_TransferLiteralToInfeedFn( manager_, tpu_executor->se_executor(), &c_literal, status.c_status); + ApiConverter::Free(&c_literal); + return status.status(); } @@ -135,6 +137,9 @@ Status TpuTransferManager::TransferLiteralFromOutfeed( manager_, tpu_executor->se_executor(), &c_shape, &c_literal, status.c_status); + ApiConverter::Free(&c_shape); + ApiConverter::Free(&c_literal); + return status.status(); } @@ -265,6 +270,7 @@ Status TpuTransferManager::LinearizeToBuffers( tpu::ExecutorApiFn()->TpuTransferManager_FreeBuffersFn( buffers_array, buffers_size, buffers_array_size); + ApiConverter::Free(&c_literal); return status.status(); }