diff --git a/spacy/ml/parser_model.pyx b/spacy/ml/parser_model.pyx
index 5d5b5eaf9..8bb1ad4a6 100644
--- a/spacy/ml/parser_model.pyx
+++ b/spacy/ml/parser_model.pyx
@@ -227,16 +227,16 @@ cdef int arg_max_if_valid(const weight_t* scores, const int* is_valid, int n) no
                 best = i
     return best
 
-cdef inline int _arg_max(float[:] scores, const int n_classes) nogil:
-    if n_classes == 2:
-        return 0 if scores[0] > scores[1] else 1
+cdef int _arg_max(const float* scores, const int n_classes) nogil:
     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
 
 class ParserStepModel(Model):
@@ -538,5 +538,3 @@ cdef class precompute_hiddens:
 
         return state_vector, backprop_relu
 
-        
-