2 files changed, 38 insertions, 0 deletions
diff --git a/src/armnn/optimizations/AddDebug.hpp b/src/armnn/optimizations/AddDebug.hpp
new file mode 100644
index 0000000000..60271b0d77
--- /dev/null
+++ b/src/armnn/optimizations/AddDebug.hpp
@@ -0,0 +1,37 @@
+//
+// Copyright © 2017 Arm Ltd. All rights reserved.
+// SPDX-License-Identifier: MIT
+//
+#pragma once
+
+#include "Optimization.hpp"
+#include "NetworkUtils.hpp"
+
+namespace armnn
+{
+namespace optimizations
+{
+
+class AddDebugImpl
+{
+public:
+
+    void Run(Graph& graph, Layer& layer) const
+    {
+        if (layer.GetType() != LayerType::Debug && layer.GetType() != LayerType::Output)
+        {
+            // if the inputs/outputs of this layer do not have a debug layer
+            // insert the debug layer after them
+            InsertDebugLayerAfter(graph, layer);
+        }
+    }
+
+protected:
+    AddDebugImpl() = default;
+    ~AddDebugImpl() = default;
+};
+
+using InsertDebugLayer = OptimizeForType<Layer, AddDebugImpl>;
+
+} // namespace optimizations
+} // namespace armnn
diff --git a/src/armnn/optimizations/All.hpp b/src/armnn/optimizations/All.hpp
index a1bff3c5e1..0a6684ee3b 100644
--- a/src/armnn/optimizations/All.hpp
+++ b/src/armnn/optimizations/All.hpp
@@ -12,3 +12,4 @@
 #include "MovePermuteUp.hpp"
 #include "OptimizeInverseConversions.hpp"
 #include "ConvertFp32NetworkToFp16.hpp"
+#include "AddDebug.hpp"