From 72f02b890912a3323a1d477648bac013e403cfb3 Mon Sep 17 00:00:00 2001
From: saoste01 <saoirse.stewart@arm.com>
Date: Fri, 2 Nov 2018 17:54:10 +0000
Subject: MLCE-65: Update for DepthwiseConvolution issue

Change-Id: Icad9d16c840049bf182b8784714734524fdb0878
---
 src/backends/cl/workloads/ClDepthwiseConvolutionWorkload.cpp | 5 ++++-
 1 file changed, 4 insertions(+), 1 deletion(-)

diff --git a/src/backends/cl/workloads/ClDepthwiseConvolutionWorkload.cpp b/src/backends/cl/workloads/ClDepthwiseConvolutionWorkload.cpp
index 6b159f15e4..83aa03d727 100644
--- a/src/backends/cl/workloads/ClDepthwiseConvolutionWorkload.cpp
+++ b/src/backends/cl/workloads/ClDepthwiseConvolutionWorkload.cpp
@@ -88,8 +88,11 @@ ClDepthwiseConvolutionWorkload::ClDepthwiseConvolutionWorkload(
 
     const unsigned int depthMultiplier = weightInfo.GetShape()[0];
 
+    const unsigned int widthIndex = (m_Data.m_Parameters.m_DataLayout == DataLayout::NHWC) ? 3 : 2;
+    const unsigned int heightIndex = (m_Data.m_Parameters.m_DataLayout == DataLayout::NHWC) ? 2 : 1;
+
     //Check for optimisation opportunities.
-    bool use3x3Optimisation = (weightInfo.GetShape()[3] == 3) && (weightInfo.GetShape()[2] == 3);
+    bool use3x3Optimisation = (weightInfo.GetShape()[widthIndex] == 3) && (weightInfo.GetShape()[heightIndex] == 3);
     if (use3x3Optimisation)
     {
         m_DepthwiseConvolutionLayer = std::make_unique<arm_compute::CLDepthwiseConvolutionLayer3x3>();
-- 
cgit v1.2.1