diff options
Diffstat (limited to 'src/core/NEON/kernels/NESpaceToDepthLayerKernel.cpp')
| -rw-r--r-- | src/core/NEON/kernels/NESpaceToDepthLayerKernel.cpp | 6 |
1 files changed, 3 insertions, 3 deletions
diff --git a/src/core/NEON/kernels/NESpaceToDepthLayerKernel.cpp b/src/core/NEON/kernels/NESpaceToDepthLayerKernel.cpp index 4803365013..7b56c43994 100644 --- a/src/core/NEON/kernels/NESpaceToDepthLayerKernel.cpp +++ b/src/core/NEON/kernels/NESpaceToDepthLayerKernel.cpp @@ -66,7 +66,7 @@ Status validate_arguments(const ITensorInfo *input, const ITensorInfo *output, i } // namespace NESpaceToDepthLayerKernel::NESpaceToDepthLayerKernel() - : _input(nullptr), _output(nullptr), _block_shape() + : _input(nullptr), _output(nullptr), _block_shape(),_data_layout(DataLayout::UNKNOWN) { } @@ -82,6 +82,7 @@ void NESpaceToDepthLayerKernel::configure(const ITensor *input, ITensor *output, _input = input; _block_shape = block_shape; _output = output; + _data_layout = input->info()->data_layout(); // Configure kernel window Window win = calculate_max_window(*output->info(), Steps()); @@ -100,8 +101,7 @@ void NESpaceToDepthLayerKernel::run(const Window &window, const ThreadInfo &info ARM_COMPUTE_ERROR_ON_UNCONFIGURED_KERNEL(this); ARM_COMPUTE_ERROR_ON_INVALID_SUBWINDOW(ICPPKernel::window(), window); - const DataLayout data_layout = _input->info()->data_layout(); - const int channel_idx = get_data_layout_dimension_index(data_layout, DataLayoutDimension::CHANNEL); + const int channel_idx = get_data_layout_dimension_index(_data_layout, DataLayoutDimension::CHANNEL); const int element_size = _input->info()->element_size(); const size_t channel_size = _input->info()->dimension(channel_idx); |