aboutsummaryrefslogtreecommitdiff
path: root/src/graph/nodes/BatchNormalizationLayer.cpp
diff options
context:
space:
mode:
authorIsabella Gottardi <isabella.gottardi@arm.com>2017-11-15 16:54:03 +0000
committerAnthony Barbier <anthony.barbier@arm.com>2018-11-02 16:35:24 +0000
commitd03006a6b97f3b1830acd19bcbc3720261ade64f (patch)
treefabf212460aa39b0040de7227c093041967a8db9 /src/graph/nodes/BatchNormalizationLayer.cpp
parentce58a9f8f8504c165ca4527bfd991a4029437cba (diff)
downloadComputeLibrary-d03006a6b97f3b1830acd19bcbc3720261ade64f.tar.gz
COMPMID-556 - Fix BatchNormalizationLayer (set weights target)
Change-Id: Ibc72b1fe2f135db4f6935cc2b3280c5b0c6131f4 Reviewed-on: http://mpd-gerrit.cambridge.arm.com/95926 Tested-by: Kaizen <jeremy.johnson+kaizengerrit@arm.com> Reviewed-by: Anthony Barbier <anthony.barbier@arm.com>
Diffstat (limited to 'src/graph/nodes/BatchNormalizationLayer.cpp')
-rw-r--r--src/graph/nodes/BatchNormalizationLayer.cpp6
1 files changed, 6 insertions, 0 deletions
diff --git a/src/graph/nodes/BatchNormalizationLayer.cpp b/src/graph/nodes/BatchNormalizationLayer.cpp
index 1fc83e4f69..927ce5e728 100644
--- a/src/graph/nodes/BatchNormalizationLayer.cpp
+++ b/src/graph/nodes/BatchNormalizationLayer.cpp
@@ -61,6 +61,12 @@ std::unique_ptr<arm_compute::IFunction> BatchNormalizationLayer::instantiate_nod
bool gamma_is_loaded = _gamma.tensor() != nullptr;
bool beta_is_loaded = _beta.tensor() != nullptr;
+ // Set mean, var, gamma and beta target
+ _mean.set_target(_target_hint);
+ _var.set_target(_target_hint);
+ _gamma.set_target(_target_hint);
+ _beta.set_target(_target_hint);
+
// Create node context
NodeContext node_ctx(OperationType::BatchNormalizationLayer);
node_ctx.set_target(_target_hint);