text: import torch
  import torch.nn as nn
  
  
  class CosSim(nn.Module):
      def __init__(self, nfeat, nclass, learn_cent=True):
          super(CosSim, self).__init__()
          self.nfeat = nfeat
          self.nclass = nclass
          self.learn_cent = learn_cent
  
          self.centroids = nn.Parameter(torch.randn(nclass, nfeat))
          if not learn_cent:
              self.centroids.requires_grad_(False)
  
      def forward(self, x):
          norms = torch.norm(x, p=2, dim=-1, keepdim=True)
          nfeat = torch.div(x, norms)
  
          norms_c = torch.norm(self.centroids, p=2, dim=-1, keepdim=True)
          ncenters = torch.div(self.centroids, norms_c)
          logits = torch.matmul(nfeat, torch.transpose(ncenters, 0, 1))
  
          return logits
  
      def extra_repr(self) -> str:
          return 'in_features={}, n_class={}, learn_centroid={}'.format(
              self.nfeat, self.nclass, self.learn_cent
          )
  