Update app.py

app.py

@@ -99,7 +99,7 @@ def plot_figure(N, u1, std1, u2, std2, show_dist, classifier=None):
     #zz = np.array(  [model_sk.predict( [[xx,yy]])[0] for xx, yy in zip(np.ravel(X), np.ravel(Y)) ] )
 
     #Reshaping the predicted class into the meshgrid shape
-    Z = zz.reshape(X.shape)
+    Z = zz.reshape(xx.shape)
 
     pp.pcolormesh(xx, yy, Z, cmap=cmap_light, alpha=0.2)
   elif classifier == 'QDA':
@@ -114,7 +114,7 @@ def plot_figure(N, u1, std1, u2, std2, show_dist, classifier=None):
     #zz = np.array(  [model_sk.predict( [[xx,yy]])[0] for xx, yy in zip(np.ravel(X), np.ravel(Y)) ] )
 
     #Reshaping the predicted class into the meshgrid shape
-    Z = zz.reshape(X.shape)
+    Z = zz.reshape(xx.shape)
     
     print("Z: ",Z)
     pp.pcolormesh(xx, yy, Z, cmap=cmap_light, alpha=0.2)