// make color map
n_classes = np.unique(y).shape[0]
colors = colors.split(",")
cmap = ListedColormap(colors[:n_classes])
// plot the decision surface
if dim == "2d":
After Change
n_classes = np.unique(y).shape[0]
colors = colors.split(",")
colors_gen = cycle(colors)
colors = [next(colors_gen) for c in range(n_classes)]
if dim == "2d":
y_min, y_max = X[:, 1].min() - 1, X[:, 1].max() + 1