@ptrblck
For the K-means loss, this is what I am doing. I just want to verify if its correct . Or is there any other clean way of doing this?
import torch
class KMeansClusteringLoss(torch.nn.Module):
def __init__(self):
    super(KMeansClusteringLoss,self).__init__()

def forward(self, encode_output, centroids):
    assert (encode_output.shape[1] == centroids.shape[1]),"Dimensions Mismatch"
    n = encode_output.shape[0]
    d = encode_output.shape[1]
    k = centroids.shape[0]

    z = encode_output.reshape(n,1,d)
    z = z.repeat(1,k,1)

    mu = centroids.reshape(1,k,d)
    mu = mu.repeat(n,1,1)

    dist = (z-mu).norm(2,dim=2).reshape((n,k))
    loss = dist.min(dim=1)[0].mean()

    return loss