experiments/STT-tensorflow
1
0
Fork 0
Code Issues Pull Requests Projects Releases Wiki Activity

Fix memory leaks in TpuTransferManager by calling ApiConverter::Free

Browse Source 
PiperOrigin-RevId: 326347833
Change-Id: I4d6c3935cdcc56c16842e7985f10e44c2b0c34a1
This commit is contained in:
Frank Chen 2020-08-12 17:23:42 -07:00 committed by TensorFlower Gardener
parent ac09152e88
commit 8b03b9681e
1 changed files with 6 additions and 0 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

6
tensorflow/stream_executor/tpu/tpu_transfer_manager.cc
View File

@ -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();
}