diff --git a/src/main/resources/templates/gluon/elements/OneHot.ftl b/src/main/resources/templates/gluon/elements/OneHot.ftl
index 40991a3a62075fa02b533a0138fd8d44a2bb3688..6e75def6499f5a8b184bca6e83c2870df319fa05 100644
--- a/src/main/resources/templates/gluon/elements/OneHot.ftl
+++ b/src/main/resources/templates/gluon/elements/OneHot.ftl
@@ -1,8 +1,8 @@
 <#assign input = element.inputs[0]>
 <#assign size = element.size>
 <#if mode == "ARCHITECTURE_DEFINITION">
-            self.${element.name} = OneHot(size=${size})
-            <#include "OutputShape.ftl">
+        self.${element.name} = OneHot(size=${size})
+        <#include "OutputShape.ftl">
 </#if>
 <#if mode == "FORWARD_FUNCTION">
         ${element.name} = self.${element.name}(${input})