code: def test_p2s_grad_loss_with_paper_formula(self):
      torch.manual_seed(23)
      num_classes = 20
      batch_size = 100
      descriptors_dim = 128
      for dtype in TEST_DTYPES:
          if dtype == torch.float16:
              continue
          embeddings = torch.randn(
              batch_size,
              descriptors_dim,
              device=TEST_DEVICE,
              dtype=dtype,
              requires_grad=True,
          )
          labels = torch.randint(
              0, num_classes, size=(batch_size,), device=TEST_DEVICE
          )
  
          loss_fn = P2SGradLoss(
              descriptors_dim=descriptors_dim, num_classes=num_classes
          )
          optimizer = torch.optim.SGD([loss_fn.weight], lr=0.001)
          optimizer.zero_grad()
          copy_weights = loss_fn.weight.data.clone()
  
          loss = loss_fn(embeddings, labels)
          loss.backward()
  
          copy_weights = c_f.to_device(
              copy_weights, device=TEST_DEVICE, dtype=dtype
          )
          for j in range(num_classes):
              w_j = copy_weights[:, j]
              x_norm = torch.norm(embeddings, p=2, dim=1).view(-1, 1)
              w_norm = torch.norm(w_j, p=2, dim=0)
              cos_theta_j = embeddings.mm(w_j.unsqueeze(1)) / (
                  x_norm * w_norm
              )
              cos_theta_j = cos_theta_j.view(-1, 1)
  
              index = torch.zeros(labels.shape, device=TEST_DEVICE)
              index[labels == j] = 1
              D_j = embeddings / x_norm
              L_j = cos_theta_j - index.unsqueeze(1)
  
              gradients = 2 * torch.mean(L_j * D_j, dim=0) / num_classes
              gradients = gradients.to(dtype)
              self.assertTrue(
                  torch.all(
                      torch.isclose(
                          gradients,
                          loss_fn.weight.grad[:, j],
                          rtol=0.5e-1,
                      )
                  )
              )
  
file_path: pytorch-metric-learning/tests/losses/test_p2s_grad_loss.py
imports: import unittest
  import torch
  import torch.nn
  from pytorch_metric_learning.losses import P2SGradLoss
  from pytorch_metric_learning.utils import common_functions
  from ..None import TEST_DEVICE, TEST_DTYPES
label_desc: Use image node_id 1 for calling the TestP2SGradLoss obj's underlying member method code with example usage: obj.test_p2s_grad_loss_with_paper_formula() without return types