diff --git a/torch_modules/src/ContextModule.cpp b/torch_modules/src/ContextModule.cpp
index a6e1c85974c7b97a60cc6bd0980dc3125163593f..4b973c071d4b8247760426c88977521ad2970a1e 100644
--- a/torch_modules/src/ContextModule.cpp
+++ b/torch_modules/src/ContextModule.cpp
@@ -107,7 +107,24 @@ void ContextModuleImpl::addToContext(std::vector<std::vector<long>> & context, c
       }
       else
       {
-        int dictIndex = dict.getIndexOrInsert(functions[colIndex](config.getAsFeature(col, index)));
+        int dictIndex;
+        if (col == Config::idColName)
+        {
+          std::string value;
+          if (config.isCommentPredicted(index))
+            value = "ID(comment)";
+          else if (config.isMultiwordPredicted(index))
+            value = "ID(multiword)";
+          else if (config.isTokenPredicted(index))
+            value = "ID(token)";
+          dictIndex = dict.getIndexOrInsert(value);
+        }
+        else if (col == Config::EOSColName)
+        {
+          dictIndex = dict.getIndexOrInsert(fmt::format("EOS({})", config.getAsFeature(col, index)));
+        }
+        else
+          dictIndex = dict.getIndexOrInsert(functions[colIndex](config.getAsFeature(col, index)));
 
         for (auto & contextElement : context)
           contextElement.push_back(dictIndex);