diff --git a/toolkit.py b/toolkit.py
index 0403f9bfa0183311d47082c294449130f49009af..5ac5fff1d9a9798481b9f5bd9e190a97a3abb697 100755
--- a/toolkit.py
+++ b/toolkit.py
@@ -2107,7 +2107,7 @@ class ClassificationProjectDecorr(ClassificationProject):
                 )
                 if (dropout_fraction is not None) and (dropout_fraction > 0):
                     layers.append(Dropout(rate=dropout_fraction))
-            layers.append(Dense(1, activation=self.activation_function_output))
+            layers.append(Dense(1, activation=self.activation_function_output, name="class"))
         return self._class_layers
 
 
@@ -2120,8 +2120,8 @@ class ClassificationProjectDecorr(ClassificationProject):
             self._adv_hidden_layers = []
             self._adv_target_layers = []
             self._adv_hidden_layers.append(Dense(128, activation="tanh"))
-            for binning in self.decorr_binnings:
-                layer = Dense(len(binning), activation="softmax")
+            for binning, field_name in zip(self.decorr_binnings, self.target_fields):
+                layer = Dense(len(binning), activation="softmax", name="adv_"+field_name)
                 self._adv_target_layers.append(layer)
         return self._adv_hidden_layers+self._adv_target_layers