labels = np.array([num for _ in range(n_row) for num in range(n_row)])
One of the n_row in this line should be n_classes?
