aboutsummaryrefslogtreecommitdiff
path: root/arm_compute/runtime/NEON/functions/NESoftmaxLayer.h
diff options
context:
space:
mode:
Diffstat (limited to 'arm_compute/runtime/NEON/functions/NESoftmaxLayer.h')
-rw-r--r--arm_compute/runtime/NEON/functions/NESoftmaxLayer.h4
1 files changed, 3 insertions, 1 deletions
diff --git a/arm_compute/runtime/NEON/functions/NESoftmaxLayer.h b/arm_compute/runtime/NEON/functions/NESoftmaxLayer.h
index 01402aee63..a265f70043 100644
--- a/arm_compute/runtime/NEON/functions/NESoftmaxLayer.h
+++ b/arm_compute/runtime/NEON/functions/NESoftmaxLayer.h
@@ -27,6 +27,7 @@
#include "arm_compute/core/NEON/kernels/NEFillBorderKernel.h"
#include "arm_compute/core/NEON/kernels/NESoftmaxLayerKernel.h"
#include "arm_compute/runtime/IFunction.h"
+#include "arm_compute/runtime/MemoryGroup.h"
#include "arm_compute/runtime/Tensor.h"
namespace arm_compute
@@ -47,7 +48,7 @@ class NESoftmaxLayer : public IFunction
{
public:
/** Constructor */
- NESoftmaxLayer();
+ NESoftmaxLayer(std::shared_ptr<IMemoryManager> memory_manager = nullptr);
/** Set the input and output tensors.
*
* @param[in] input Source tensor. Data types supported: QS8/QS16/F16/F32.
@@ -59,6 +60,7 @@ public:
void run() override;
private:
+ MemoryGroup _memory_group;
NELogits1DMaxKernel _max_kernel;
NELogits1DShiftExpSumKernel _shift_exp_sum_kernel;
NELogits1DNormKernel _norm_kernel;