Weight Clustering in Pytorch

Question

I have been trying to do weight clustering in PyTorch for my project. For this I have used the following very rudimentary code :

from sklearn.cluster import KMeans
# from kmeans_pytorch import kmeans, kmeans_predict

'''Going to go through all the layers --> obtain their weights--> use scikit learn to do the clustering--> replace these weight by their centroids'''
def weight_clustering(model):
  model.to('cpu')
  with torch.no_grad():
    for name, params in model.named_parameters():
      param_shape=list(params.size())
      # print("the shape is ",param_shape)
      weights=torch.flatten(params)
      weights=params.reshape(-1,1)
      kmeans = KMeans(n_clusters=3, random_state=0).fit(weights)

      cluster_centers=torch.from_numpy(kmeans.cluster_centers_)
      for i in range(0,len(kmeans.labels_)):
         if kmeans.labels_[i]==0:
          weights[i]=cluster_centers[0]
         elif kmeans.labels_[i]==1:
          weights[i]=cluster_centers[1]
         elif kmeans.labels_[i]==2:
          weights[i]=cluster_centers[2]
         elif kmeans.labels_[i]==3:
          weights[i]=cluster_centers[3]
         else :
          weights[i]=cluster_centers[4]
         

      reshape_size_tuple=tuple(param_shape)
      weights=weights.reshape(reshape_size_tuple)
      params.data=weights.data

However, this would save the weights individually and hence would not achieve any kind of reduction in the model size. I wanted to use and access the same storage (something similar to the following link : PyTorch Views ). I believe to the best of my knowledge that PyTorch does not have any clustering mechanism as yet as opposed to TensorFlow.

I have tried using view in the above code but have not been able to do so successfully getting the desired results.Can someone guide/help me with this?