diff --git a/src/Dialect/ONNX/ElementsAttr/DisposablePool.cpp b/src/Dialect/ONNX/ElementsAttr/DisposablePool.cpp index aa0f24d244..bdc654f540 100644 --- a/src/Dialect/ONNX/ElementsAttr/DisposablePool.cpp +++ b/src/Dialect/ONNX/ElementsAttr/DisposablePool.cpp @@ -114,7 +114,7 @@ void DisposablePool::scrub(ModuleOp moduleOp, OpAttrDictionary opsAttrs) { return llvm::make_range(batchBegin, batchEnd); }; // Parallel worker body: Fetch and process batches until there are no more. - auto work = [&fetchBatch](size_t threadNumber) { + auto work = [&fetchBatch, &translationMutex](size_t threadNumber) { for (;;) { auto batch = fetchBatch(); if (batch.empty()) @@ -122,9 +122,10 @@ void DisposablePool::scrub(ModuleOp moduleOp, OpAttrDictionary opsAttrs) { for (auto &[id, translation] : batch) { auto &[disposable, dense] = translation; dense = disposable.toDenseElementsAttr(); - // TODO: Consider calling disposable.dispose() here to free up memory - // on the go to make memory available to create the next - // DenseElementsAttr. In that case should we lock mutex? + { + const std::lock_guard lock(translationMutex); + disposable.dispose(); + } } } };