diff --git a/spacy/ml/parser_model.pyx b/spacy/ml/parser_model.pyx
index 7375375a7..8698fc5b6 100644
--- a/spacy/ml/parser_model.pyx
+++ b/spacy/ml/parser_model.pyx
@@ -527,4 +527,14 @@ cdef class precompute_hiddens:
 
         return state_vector, backprop_relu
 
-cdef inline int _arg_max(const float* scores, const int n_classes) nogil:
+    cdef inline int _arg_max(const float* scores, const int n_classes) nogil:
+        if n_classes == 2:
+            return 0 if scores[0] > scores[1] else 1
+        cdef int i
+        cdef int best = 0
+        cdef float mode = scores[0]
+        for i in range(1, n_classes):
+            if scores[i] > mode:
+                mode = scores[i]
+                best = i
+        return best