text: import torch
  import torch.nn as nn
  import torch.nn.functional as F
  import math
  
  
  class ArcMarginProduct(nn.Module):
      def __init__(self, in_features, out_features):
          super().__init__()
          self.weight = nn.Parameter(torch.FloatTensor(out_features, in_features))
          self.reset_parameters()
  
      def reset_parameters(self):
          stdv = 1. / math.sqrt(self.weight.size(1))
          self.weight.data.uniform_(-stdv, stdv)
  
      def forward(self, features):
          cosine = F.linear(F.normalize(features), F.normalize(self.weight))
          return cosine