Does not work with 2D, 3D features

[morawi]

How to set the num_features?

[morawi]

I gave the code a go, since I was iterating over K=1,..., n; seems that K=1 was the source of the problem.
Although there might not be sense of using K=1, we can prevent the error using:

Line 57 
dimens = 0 if num_clusters==1 else 1
choice_cluster = torch.argmin(dis, dim=dimens)   # choice_cluster = torch.argmin(dis, dim=1) 

dim=dimens should replace all dim=1 instances

Or, adding throw error message whenever num_clusters <2