def forward(self, embedding, y_true):
    cos_theta = F.linear(F.normalize(embedding), F.normalize(self.weight)).clamp(((- 1) + self.eps), (1 - self.eps))
    if (y_true.dim() != 1):
        y_true_one_hot = y_true.float()
    else:
        y_true_one_hot = torch.zeros_like(cos_theta)
        y_true_one_hot.scatter_(1, y_true.unsqueeze(1), 1.0)
    with torch.no_grad():
        B_batch = cos_theta[y_true_one_hot.eq(0)].mul(self.prev_s).exp().sum().div(embedding.size(0))
        self.running_B = ((self.running_B * self.momentum) + (B_batch * (1 - self.momentum)))
        theta = torch.acos(cos_theta.clamp(((- 1) + self.eps), (1 - self.eps)))
        theta_batch = theta[y_true_one_hot.ne(0)].mean().clamp_max((math.pi / 4))
        self.running_theta = ((self.running_theta * self.momentum) + (theta_batch * (1 - self.momentum)))
        self.prev_s = (self.running_B.log() / torch.cos(self.running_theta))
    self.idx += 1
    if ((self.idx % 1000) == 0):
        print(f'''
Running B: {self.running_B:.2f}. Running theta: {self.running_theta:.2f}. Running S: {self.prev_s:.2f}''')
    return self.final_criterion((cos_theta * self.prev_s), y_true_one_hot)
def forward(self, embedding, y_true):
    cos_theta = F.linear(F.normalize(embedding), F.normalize(self.weight)).clamp(((- 1) + self.eps), (1 - self.eps))
    if (y_true.dim() != 1):
        y_true_one_hot = y_true.float()
    else:
        y_true_one_hot = torch.zeros_like(cos_theta)
        y_true_one_hot.scatter_(1, y_true.unsqueeze(1), 1.0)
    with torch.no_grad():
        B_batch = cos_theta[y_true_one_hot.eq(0)].mul(self.prev_s).exp().sum().div(embedding.size(0))
        self.running_B = ((self.running_B * self.momentum) + (B_batch * (1 - self.momentum)))
        theta = torch.acos(cos_theta.clamp(((- 1) + self.eps), (1 - self.eps)))
        theta_batch = theta[y_true_one_hot.ne(0)].mean().clamp_max((math.pi / 4))
        self.running_theta = ((self.running_theta * self.momentum) + (theta_batch * (1 - self.momentum)))
        self.prev_s = (self.running_B.log() / torch.cos(self.running_theta))
    self.idx += 1
    if ((self.idx % 1000) == 0):
        print(f'
Running B: {self.running_B:.2f}. Running theta: {self.running_theta:.2f}. Running S: {self.prev_s:.2f}')
    return self.final_criterion((cos_theta * self.prev_s), y_true_one_hot)