From 8b03b9681e605bb55ac3605d3b4931190fe98add Mon Sep 17 00:00:00 2001 From: Frank Chen Date: Wed, 12 Aug 2020 17:23:42 -0700 Subject: [PATCH] Fix memory leaks in TpuTransferManager by calling ApiConverter::Free PiperOrigin-RevId: 326347833 Change-Id: I4d6c3935cdcc56c16842e7985f10e44c2b0c34a1 --- tensorflow/stream_executor/tpu/tpu_transfer_manager.cc | 6 ++++++ 1 file changed, 6 insertions(+) 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(); }