1 files changed, 9 insertions, 2 deletions
diff --git a/arm_compute/runtime/CL/functions/CLFlattenLayer.h b/arm_compute/runtime/CL/functions/CLFlattenLayer.h
index f5f4ff554f..ffe06aa610 100644
--- a/arm_compute/runtime/CL/functions/CLFlattenLayer.h
+++ b/arm_compute/runtime/CL/functions/CLFlattenLayer.h
@@ -25,7 +25,8 @@
 #define ARM_COMPUTE_CLFLATTENLAYER_H
 
 #include "arm_compute/core/Types.h"
-#include "arm_compute/runtime/CL/ICLSimpleFunction.h"
+#include "arm_compute/runtime/CL/functions/CLReshapeLayer.h"
+#include "arm_compute/runtime/IFunction.h"
 
 namespace arm_compute
 {
@@ -38,7 +39,7 @@ class ITensorInfo;
 * -# @ref CLFlattenLayerKernel
 *
 */
-class CLFlattenLayer : public ICLSimpleFunction
+class CLFlattenLayer : public IFunction
 {
 public:
     /** Initialise the kernel's input and output.
@@ -68,6 +69,12 @@ public:
      * @return a status
      */
     static Status validate(const ITensorInfo *input, const ITensorInfo *output);
+
+    // Inherited methods overridden:
+    void run() override;
+
+private:
+    CLReshapeLayer _reshape{};
 };
 } // namespace arm_compute