A more easily implenment of center loss.
class CenterLoss(nn.Module):

    def __init__(self,center_num,feature_dim):
        super(CenterLoss,self).__init__()
        self.center_num = center_num
        self.feature_dim = feature_dim
        self.center_features = nn.Parameter(torch.Tensor(self.center_num,self.feature_dim))
        nn.init.normal(self.center_features.data,mean=0,std=0.1)

    def forward(self,x,label):
        B = x.size()[0]
        center = torch.index_select(self.center_features,0,label)
        diff = x-center
        loss = diff.pow(2).sum() / B
        return loss