diff --git a/src/backends/neon/workloads/NeonBatchNormalizationWorkload.hpp b/src/backends/neon/workloads/NeonBatchNormalizationWorkload.hpp
index 52e4db7..3619ea0 100644
--- a/src/backends/neon/workloads/NeonBatchNormalizationWorkload.hpp
+++ b/src/backends/neon/workloads/NeonBatchNormalizationWorkload.hpp
@@ -5,7 +5,12 @@
 
 #pragma once
 
-#include <neon/workloads/NeonWorkloadUtils.hpp>
+#include <backendsCommon/Workload.hpp>
+
+#include <arm_compute/runtime/IFunction.h>
+#include <arm_compute/runtime/Tensor.h>
+
+#include <memory>
 
 namespace armnn
 {
@@ -26,7 +31,7 @@
     virtual void Execute() const override;
 
 private:
-    mutable arm_compute::NEBatchNormalizationLayer m_Layer;
+    std::unique_ptr<arm_compute::IFunction> m_Layer;
 
     std::unique_ptr<arm_compute::Tensor> m_Mean;
     std::unique_ptr<arm_compute::Tensor> m_Variance;