diff --git a/torch_modules/src/ContextModule.cpp b/torch_modules/src/ContextModule.cpp
index a3129c8eef06a2cf41056571c8b50eb95180f59d..48240340b0718c2ef72cf0fa50bbb60b219e00da 100644
--- a/torch_modules/src/ContextModule.cpp
+++ b/torch_modules/src/ContextModule.cpp
@@ -113,7 +113,7 @@ void ContextModuleImpl::addToContext(std::vector<std::vector<long>> & context, c
       if (index == -1)
       {
         for (auto & contextElement : context)
-          contextElement.push_back(dict.getIndexOrInsert(Dict::nullValueStr));
+          contextElement.push_back(dict.getIndexOrInsert(fmt::format("{}({})", col, Dict::nullValueStr)));
       }
       else
       {
@@ -134,7 +134,12 @@ void ContextModuleImpl::addToContext(std::vector<std::vector<long>> & context, c
           dictIndex = dict.getIndexOrInsert(fmt::format("EOS({})", config.getAsFeature(col, index)));
         }
         else
-          dictIndex = dict.getIndexOrInsert(functions[colIndex](config.getAsFeature(col, index)));
+        {
+          std::string featureValue = functions[colIndex](config.getAsFeature(col, index));
+          if (w2vFile.empty())
+            featureValue = fmt::format("{}({})", col, featureValue);
+          dictIndex = dict.getIndexOrInsert(featureValue);
+        }
 
         for (auto & contextElement : context)
           contextElement.push_back(dictIndex);