diff --git a/torch_modules/src/DepthLayerTreeEmbedding.cpp b/torch_modules/src/DepthLayerTreeEmbedding.cpp
index f21596038c3eb4c87527d0b47f79a8b272e5b864..6e1342a0270a0dea0fcc4d41bef758a050709595 100644
--- a/torch_modules/src/DepthLayerTreeEmbedding.cpp
+++ b/torch_modules/src/DepthLayerTreeEmbedding.cpp
@@ -70,13 +70,7 @@ void DepthLayerTreeEmbeddingImpl::addToContext(std::vector<std::vector<long>> &
             newChilds.insert(newChilds.end(), val.begin(), val.end());
           }
         childs = newChilds;
-        if (depth == 0)
-        {
-          newChilds.clear();
-          auto gov = config.has(0,index,0) ? config.getAsFeature(Config::headColName, index).get() : "-1";
-          newChilds.emplace_back(util::isEmpty(gov) ? "-1" : gov);
-          newChilds.insert(newChilds.end(), childs.begin(), childs.end());
-        }
+
         for (int i = 0; i < maxElemPerDepth[depth]; i++)
           for (auto & col : columns)
             if (i < (int)newChilds.size() and config.has(col, std::stoi(newChilds[i]), 0))
diff --git a/torch_modules/src/LSTMNetwork.cpp b/torch_modules/src/LSTMNetwork.cpp
index 7a25024b67d90bd2f5f2acde0839baa308e4bc74..bf776262b90258613d34a627dcbda2cf47184553 100644
--- a/torch_modules/src/LSTMNetwork.cpp
+++ b/torch_modules/src/LSTMNetwork.cpp
@@ -26,7 +26,7 @@ LSTMNetworkImpl::LSTMNetworkImpl(int nbOutputs, int unknownValueThreshold, std::
   if (!treeEmbeddingColumns.empty())
   {
     hasTreeEmbedding = true;
-    treeEmbedding = register_module("treeEmbedding", DepthLayerTreeEmbedding(treeEmbeddingNbElems,embeddingsSize,treeEmbeddingSize,treeEmbeddingColumns,treeEmbeddingBuffer,treeEmbeddingStack,lstmOptionsAll));
+    treeEmbedding = register_module("treeEmbedding", DepthLayerTreeEmbedding(treeEmbeddingNbElems,embeddingsSize,treeEmbeddingSize,treeEmbeddingColumns,treeEmbeddingBuffer,treeEmbeddingStack,lstmOptions));
     treeEmbedding->setFirstInputIndex(currentInputSize);
     currentOutputSize += treeEmbedding->getOutputSize();
     currentInputSize += treeEmbedding->getInputSize();