From d03006a6b97f3b1830acd19bcbc3720261ade64f Mon Sep 17 00:00:00 2001
From: Isabella Gottardi <isabella.gottardi@arm.com>
Date: Wed, 15 Nov 2017 16:54:03 +0000
Subject: 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>
---
 src/graph/nodes/BatchNormalizationLayer.cpp | 6 ++++++
 1 file changed, 6 insertions(+)

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);
-- 
cgit v1.2.1