2 files changed, 2 insertions, 4 deletions
diff --git a/tests/ImageTensorGenerator/ImageTensorGenerator.cpp b/tests/ImageTensorGenerator/ImageTensorGenerator.cpp
index b4432558c4..a69a098eb4 100644
--- a/tests/ImageTensorGenerator/ImageTensorGenerator.cpp
+++ b/tests/ImageTensorGenerator/ImageTensorGenerator.cpp
@@ -9,7 +9,6 @@
 #include <armnn/TypesUtils.hpp>
 #include <armnnUtils/Filesystem.hpp>
 
-#include <mapbox/variant.hpp>
 #include <cxxopts/cxxopts.hpp>
 
 #include <algorithm>
@@ -297,9 +296,7 @@ int main(int argc, char* argv[])
     const unsigned int batchSize = 1;
     const armnn::DataLayout outputLayout(cmdline.GetLayout());
 
-    using TContainer = mapbox::util::variant<std::vector<float>, std::vector<int>, std::vector<uint8_t>,
-                                             std::vector<int8_t>>;
-    std::vector<TContainer> imageDataContainers;
+    std::vector<armnn::TContainer> imageDataContainers;
     const NormalizationParameters& normParams = GetNormalizationParameters(modelFormat, outputType);
     try
     {
diff --git a/tests/ImageTensorGenerator/ImageTensorGenerator.hpp b/tests/ImageTensorGenerator/ImageTensorGenerator.hpp
index 6d2e549360..c668608f1d 100644
--- a/tests/ImageTensorGenerator/ImageTensorGenerator.hpp
+++ b/tests/ImageTensorGenerator/ImageTensorGenerator.hpp
@@ -5,6 +5,7 @@
 
 #include "../InferenceTestImage.hpp"
 
+#include <armnn/Utils.hpp>
 #include <armnn/TypesUtils.hpp>
 
 #include <armnnUtils/Permute.hpp>