1 files changed, 13 insertions, 3 deletions
diff --git a/src/backends/reference/workloads/RefReshapeWorkload.cpp b/src/backends/reference/workloads/RefReshapeWorkload.cpp
index 6d29781937..960d591fec 100644
--- a/src/backends/reference/workloads/RefReshapeWorkload.cpp
+++ b/src/backends/reference/workloads/RefReshapeWorkload.cpp
@@ -14,11 +14,21 @@ namespace armnn
 
 void RefReshapeWorkload::Execute() const
 {
+    Execute(m_Data.m_Inputs, m_Data.m_Outputs);
+}
+
+void RefReshapeWorkload::ExecuteAsync(WorkingMemDescriptor &workingMemDescriptor)
+{
+    Execute(workingMemDescriptor.m_Inputs, workingMemDescriptor.m_Outputs);
+}
+
+void RefReshapeWorkload::Execute(std::vector<ITensorHandle*> inputs, std::vector<ITensorHandle*> outputs) const
+{
     ARMNN_SCOPED_PROFILING_EVENT(Compute::CpuRef, "RefReshapeWorkload_Execute");
 
-    void* output = GetOutputTensorData<void>(0, m_Data);
-    const void* input = GetInputTensorData<void>(0, m_Data);
-    unsigned int numBytes = GetTensorInfo(m_Data.m_Inputs[0]).GetNumBytes();
+    void* output =  outputs[0]->Map();
+    const void* input =  inputs[0]->Map();
+    unsigned int numBytes = GetTensorInfo(inputs[0]).GetNumBytes();
     memcpy(output, input, numBytes);
 }