1 files changed, 7 insertions, 8 deletions

diff --git a/src/runtime/NEON/functions/NEConcatenateLayer.cpp b/src/runtime/NEON/functions/NEConcatenateLayer.cpp
index fa7b91c3ca..e02c0c2c7a 100644
--- a/src/runtime/NEON/functions/NEConcatenateLayer.cpp
+++ b/src/runtime/NEON/functions/NEConcatenateLayer.cpp
@@ -44,10 +44,10 @@ NEConcatenateLayer::NEConcatenateLayer()
 {
 }
 
-void NEConcatenateLayer::configure(const std::vector<ITensor *> &inputs_vector, ITensor *output, DataLayoutDimension axis)
+void NEConcatenateLayer::configure(const std::vector<ITensor *> &inputs_vector, ITensor *output, size_t axis)
 {
     ARM_COMPUTE_ERROR_ON(output == nullptr);
-    _axis       = get_data_layout_dimension_index(output->info()->data_layout(), axis);
+    _axis       = axis;
     _num_inputs = inputs_vector.size();
 
     std::vector<ITensorInfo *> inputs_vector_info;
@@ -104,22 +104,21 @@ void NEConcatenateLayer::configure(const std::vector<ITensor *> &inputs_vector,
     }
 }
 
-Status NEConcatenateLayer::validate(const std::vector<ITensorInfo *> &inputs_vector, const ITensorInfo *output, DataLayoutDimension axis)
+Status NEConcatenateLayer::validate(const std::vector<ITensorInfo *> &inputs_vector, const ITensorInfo *output, size_t axis)
 {
     ARM_COMPUTE_RETURN_ERROR_ON_NULLPTR(output);
     ARM_COMPUTE_RETURN_ERROR_ON(inputs_vector.size() < 2);
-    const unsigned int _axis = get_data_layout_dimension_index(inputs_vector[0]->data_layout(), axis);
 
     // Output auto inizialitation if not yet initialized
     TensorInfo  tmp_output_info = *output->clone();
     TensorShape output_shape{};
-    if(_axis == Window::DimZ)
+    if(axis == Window::DimZ)
     {
         output_shape = arm_compute::misc::shape_calculator::calculate_depth_concatenate_shape(inputs_vector);
     }
     else
     {
-        output_shape = arm_compute::misc::shape_calculator::calculate_concatenate_shape(inputs_vector, _axis);
+        output_shape = arm_compute::misc::shape_calculator::calculate_concatenate_shape(inputs_vector, axis);
     }
     auto_init_if_empty(tmp_output_info, output_shape, 1, inputs_vector[0]->data_type());
 
@@ -127,7 +126,7 @@ Status NEConcatenateLayer::validate(const std::vector<ITensorInfo *> &inputs_vec
     for(const auto &input : inputs_vector)
     {
         ARM_COMPUTE_RETURN_ERROR_ON_NULLPTR(input);
-        switch(_axis)
+        switch(axis)
         {
             case Window::DimX:
             {
@@ -147,7 +146,7 @@ Status NEConcatenateLayer::validate(const std::vector<ITensorInfo *> &inputs_vec
             default:
                 ARM_COMPUTE_ERROR("Axis not supported");
         }
-        offset += input->dimension(_axis);
+        offset += input->dimension(axis);
     }
 
     return Status{};