1 files changed, 6 insertions, 3 deletions
diff --git a/src/backends/reference/workloads/SpaceToBatchNd.hpp b/src/backends/reference/workloads/SpaceToBatchNd.hpp
index e74e457a72..57c9b6bc25 100644
--- a/src/backends/reference/workloads/SpaceToBatchNd.hpp
+++ b/src/backends/reference/workloads/SpaceToBatchNd.hpp
@@ -5,17 +5,20 @@
 
 #pragma once
 
+#include "BaseIterator.hpp"
+#include "Decoders.hpp"
+#include "Encoders.hpp"
+
 #include <armnn/Descriptors.hpp>
 #include "armnn/Tensor.hpp"
 
 namespace armnn
 {
 
-template <typename T>
 void SpaceToBatchNd(const TensorInfo& inputInfo,
                     const TensorInfo& outputInfo,
                     const SpaceToBatchNdDescriptor& params,
-                    const T* inputData,
-                    T* outputData);
+                    Decoder<float>& inputData,
+                    Encoder<float>& outputData);
 
 } //namespace armnn