diff --git a/psiflow/learning.py b/psiflow/learning.py
index 1a5b0b7..a9e867d 100644
--- a/psiflow/learning.py
+++ b/psiflow/learning.py
@@ -86,7 +86,7 @@ def evaluate_outputs(
             isnan(errors[i]),
         )
 
-        _ = assign_identifier(state, identifier, error_discard)
+        _ = assign_identifier(state, identifier, reset)
         assigned = unpack_i(_, 0)
         identifier = unpack_i(_, 1)
         processed_states.append(assigned)