1 files changed, 8 insertions, 3 deletions
diff --git a/src/backends/neon/workloads/NeonSoftmaxUint8Workload.cpp b/src/backends/neon/workloads/NeonSoftmaxUint8Workload.cpp
index cff869c9b7..f780589075 100644
--- a/src/backends/neon/workloads/NeonSoftmaxUint8Workload.cpp
+++ b/src/backends/neon/workloads/NeonSoftmaxUint8Workload.cpp
@@ -5,6 +5,10 @@
 
 #include "NeonSoftmaxUint8Workload.hpp"
 
+#include "NeonWorkloadUtils.hpp"
+
+#include <arm_compute/runtime/NEON/functions/NESoftmaxLayer.h>
+
 namespace armnn
 {
 
@@ -12,7 +16,6 @@ NeonSoftmaxUint8Workload::NeonSoftmaxUint8Workload(const SoftmaxQueueDescriptor&
                                                    const WorkloadInfo& info,
                                                    std::shared_ptr<arm_compute::MemoryManagerOnDemand>& memoryManager)
     : Uint8Workload<SoftmaxQueueDescriptor>(descriptor, info)
-    , m_SoftmaxLayer(memoryManager)
 {
     m_Data.ValidateInputsOutputs("NeonSoftmaxUint8Workload", 1, 1);
 
@@ -27,14 +30,16 @@ NeonSoftmaxUint8Workload::NeonSoftmaxUint8Workload(const SoftmaxQueueDescriptor&
             "Invalid quantization for output. Only scale = 1.0f / 256.0f and offset = 0 supported");
     }
 
-    m_SoftmaxLayer.configure(&input, &output, descriptor.m_Parameters.m_Beta);
+    auto layer = std::make_unique<arm_compute::NESoftmaxLayer>(memoryManager);
+    layer->configure(&input, &output, descriptor.m_Parameters.m_Beta);
+    m_SoftmaxLayer.reset(layer.release());
 }
 
 void NeonSoftmaxUint8Workload::Execute() const
 {
     ARMNN_SCOPED_PROFILING_EVENT_NEON("NeonSoftmaxUint8Workload_Execute");
 
-    m_SoftmaxLayer.run();
+    m_SoftmaxLayer->run();
 }
 
 } //namespace armnn