diff options
-rw-r--r-- | src/runtime/IWeightsManager.cpp | 21 | ||||
-rw-r--r-- | src/runtime/NEON/functions/NEFullyConnectedLayer.cpp | 3 |
2 files changed, 22 insertions, 2 deletions
diff --git a/src/runtime/IWeightsManager.cpp b/src/runtime/IWeightsManager.cpp index 6dfb925fe6..b367b5f70b 100644 --- a/src/runtime/IWeightsManager.cpp +++ b/src/runtime/IWeightsManager.cpp @@ -85,6 +85,27 @@ ITensor *IWeightsManager::run(const ITensor *weights, ITransformWeights *weights } } + // Check top level weights. If all the transformations are done + // mark the weights as unused + if(_managed_weights_parents.find(weights) == _managed_weights_parents.end()) + { + auto item = _managed_weights.find(weights); + bool mark_as_unused = true; + for(auto it : item->second) + { + if(!it->is_reshape_run()) + { + mark_as_unused = false; + break; + } + } + + if(mark_as_unused) + { + weights->mark_as_unused(); + } + } + return weights_tensor; } diff --git a/src/runtime/NEON/functions/NEFullyConnectedLayer.cpp b/src/runtime/NEON/functions/NEFullyConnectedLayer.cpp index 7adc3bca9e..ee622f4699 100644 --- a/src/runtime/NEON/functions/NEFullyConnectedLayer.cpp +++ b/src/runtime/NEON/functions/NEFullyConnectedLayer.cpp @@ -426,7 +426,6 @@ void NEFullyConnectedLayer::prepare() { if(_weights_manager && _weights_manager->are_weights_managed(_original_weights)) { - cur_weights->mark_as_unused(); cur_weights = _weights_manager->run(cur_weights, &_reshape_weights_managed_function); } else @@ -455,9 +454,9 @@ void NEFullyConnectedLayer::prepare() { _converted_weights_output.allocator()->allocate(); _convert_weights.run(); + cur_weights->mark_as_unused(); } - cur_weights->mark_as_unused(); _are_weights_converted = true; } |