K-Means on Tensors

I wish to perform K-Means clustering on different datasets like MNIST,CIFAR etc. What is the recommended way to do so?
a) Convert the latent representation from tensors to numpy array and use sklearn
b) Implement k-means for tensor data in pytorch

What would be more efficient in case of CNN.