code: def test_batched_distance(self):
      torch.manual_seed(1432)
      embedding_size = 128
      for dataset_size in [15, 149, 1501, 3199, 3200]:
          # test on cpu because gpu is unreliable
          embeddings = torch.randn(
              dataset_size, embedding_size, dtype=torch.float64
          )
          ref = torch.randn(
              dataset_size, embedding_size, dtype=torch.float64
          )
          for batch_size in [32, 33, 99, 121, 128]:
              for dist_fn in [CosineSimilarity()]:
                  for use_ref in [False, True]:
                      mat = []
                      distance = BatchedDistance(
                          dist_fn, collect_fn(mat), batch_size
                      )
                      if use_ref:
                          distance(embeddings, ref)
                          correct_mat = dist_fn(embeddings, ref)
                      else:
                          distance(embeddings)
                          correct_mat = dist_fn(embeddings)
                      mat = torch.cat(mat, dim=0)
                      self.assertTrue(torch.allclose(mat, correct_mat))
  
imports: import unittest
  import torch
  from pytorch_metric_learning.distances import BatchedDistance, CosineSimilarity
label_desc: Use image node_id 1 for calling the TestBatchedDistance obj's underlying member method code with example usage: obj.test_batched_distance() without return types