Code for 07.03-05_K-means_featurization

[olutobi01]

I refer you to the fit method of the KMeansFeaturizer class for the code listing mentioned in the subject.
When the pretrain approach was used (i.e. including information from the class labels), the second re run of the k-means algorithm to get the clusters in the original space was initialized with the init argument as init=km_model_pretrain.cluster_centers_[:,:2].
I think this will only work with a two feature dataset and will break for a dataset with additional features.
Can you please clarify your approach.
Thanks