If you use triple backticks (```python) before and just the backtics (```) after your code, it will be well-formatted.
In Jupyter:
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]**2).mean() # note that you didn't have the square loss as the equation

        return loss

def sq_loss_clusters(encode_output, centroids):
    assert encode_output.size(1) == centroids.size(1), "Dimension mismatch"
    return ((encode_output[:, None]-centroids[None])**2).sum(2).min(1)[0].mean()
loss_fn = KMeansClusteringLoss()
pts = torch.randn(2000,10)
means = torch.randn(10,10)
l1 = loss_fn(pts, means)
l2 = sq_loss_clusters(pts, means)
assert (l1==l2).all()
%timeit loss_fn(pts, means)
%timeit sq_loss_clusters(pts, means)

has 3.07ms for your version and 684µs for mine, so it’s >4 faster.
I think it’s more clear what’s going on, too, if you write it more concise (maybe add a comment linking to the equation and commenting that s_ik == 1 if the assignment of point i to cluster k is optimal given fixed centers, or spread it over several lines to be better able to follow - tastes vary and that is fine).
Note that the equation you posted had the square distance, so I used that.
Best regards
Thomas