diff options
Diffstat (limited to 'src/runtime/CL/functions/CLFullyConnectedLayer.cpp')
-rw-r--r-- | src/runtime/CL/functions/CLFullyConnectedLayer.cpp | 8 |
1 files changed, 5 insertions, 3 deletions
diff --git a/src/runtime/CL/functions/CLFullyConnectedLayer.cpp b/src/runtime/CL/functions/CLFullyConnectedLayer.cpp index 6836c12022..4f9759c590 100644 --- a/src/runtime/CL/functions/CLFullyConnectedLayer.cpp +++ b/src/runtime/CL/functions/CLFullyConnectedLayer.cpp @@ -82,7 +82,7 @@ void CLFullyConnectedLayer::configure(const CLCompileContext &compile_context, c if(_impl->weights_manager != nullptr) { - _impl->weights_manager->manage(weights); + _impl->weights_manager->manage(_impl->original_weights); } if(!_impl->is_prepared) @@ -125,11 +125,13 @@ void CLFullyConnectedLayer::prepare() // Handle weights managed infrastructure if(_impl->weights_manager != nullptr && _impl->weights_manager->are_weights_managed(_impl->original_weights)) { - // If function marks b as unused ensure that all prepare stages are done before releasing + // Ensure that b gets marked as unused (memory released) only after the last function which uses b also finishes its prepare + // This is for cases where multiple functions share the same b (weights) + // Therefore when a function marks original b as unused, we pre-mark it in weights manager, and mark it back to used so that it doesn't get released before its last reference const ITensor *original_b = _impl->original_weights; if(!original_b->is_used()) { - _impl->weights_manager->mark_as_unused(original_b); + _impl->weights_manager->pre_mark_as_unused(original_b); } _impl->original_weights->mark_as_used(); _impl->weights_manager->release(_impl->original_weights); |