Decision boundary shape error

Hi,

I’m trying to plot the contour of the decision boundary using the following piece of code

# generate mesh grid of X[:,0] and X[:,1]

x_min, x_max = X_train[:, 0].min()- 0.5, X_train[:, 0].max()+0.5
y_min, y_max = X_train[:, 1].min()- 0.5, X_train[:, 1].max()+0.5

# grid spacing h

h = max(x_max-x_min,y_max-y_min)/100

xx,yy = np.meshgrid(np.arange(x_min,x_max,h),
                    np.arange(y_min,y_max,h)
                    )
r1,r2 = xx.ravel(), yy.ravel()
r1, r2 = r1.reshape((len(r1), 1)), r2.reshape((len(r2), 1))
grid = np.hstack((r1,r2))

When I pass this grid to model.predict, the shape of the output is (5858, 4). I am not able to recast it to xx.shape. Can you please let me know what I am missing?

Regards,
Adya

Nevermind, I realized I need to take the argmax of the prediction.