>>> import numpy as np
>>> a = np.array([[1,2,3],[4,3,1]])
>>> i,j = np.unravel_index(a.argmax(), a.shape)
>>> a[i,j]
4

As the other answers mentioned, you have a 2D array, so you end up with two indices. Since the array is just a row, the first index is always zero. You can bypass this in a number of ways:

1. Use prediction.argmax(). The default axis argument is None, which means operate on a flattened array. Other options that will get you the same result are prediction.argmax(-1) (last axis) and prediction.argmax(1) (second axis). Keep in mind that you will only ever get the index of the first maximum this way. That's fine if you only ever expect to have one, or only need one.

2. Use np.flatnonzero to get the linear indices similarly to the way you were doing:

   np.flatnonzero(perdiction == prediction.max())
   

3. Use np.nonzero or np.where, but extract the axis you care about:

   np.nonzero(prediction == prediction.max())[1]
   

4. ravel the array on input:

   np.where(prediction.ravel() == prediction.max())
   

5. Do the same thing, but with np.squeeze:

   np.nonzero(prediction.squeeze() == prediction.max())
   

You should use np.where

In [17]: a=np.random.uniform(0, 10, size=10)

In [18]: a
Out[18]: 
array([ 1.43249468,  4.93950873,  7.22094395,  1.20248629,  4.66783985,
        6.17578054,  4.6542771 ,  7.09244492,  7.58580515,  5.72501954])

In [20]: np.where(a==a.max())
Out[20]: (array([8]),)

This also works for 2 arrays, the returned value, is the index. Here we create a range from 1 to 9:

 x = np.arange(9.).reshape(3, 3)

This returns the index, of the the items that equal 5:

In [34]: np.where(x == 5)
Out[34]: (array([1]), array([2])) # the first one is the row index, the second is the column

You can use this value directly to slice your array:

In [35]: x[np.where(x == 5)]
Out[35]: array([ 5.])