code: def test_nearest_neighbors(self):
      # Test input with shape (batch_size, map_size_0, num_input_maps)
      input_maps = torch.tensor(
          [
              [[1, 4], [2, 5], [3, 6]],
              [[7, 10], [8, 11], [9, 12]],
              [[13, 16], [14, 17], [15, 18]],
          ]
      )
      print(input_maps.shape)
      # Test candidate maps with shape (batch_size, map_size_1, num_candidate_maps)
      candidate_maps = torch.tensor(
          [
              [[1, 1], [2, 2], [3, 3]],
              [[1, 1], [2, 2], [3, 3]],
              [[1, 1], [2, 2], [3, 3]],
          ]
      )
      print(candidate_maps.shape)
      # Test distances with shape (batch_size, map_size_0, map_size_1)
      distances = torch.tensor(
          [
              [[0, 1, 2], [1, 0, 3]],
              [[4, 3, 2], [3, 2, 1]],
              [[2, 3, 4], [3, 4, 5]],
          ]
      )
      print(input_maps.shape)
      # Test num_matches = 2
      input_maps_filtered, candidate_maps_filtered = nearest_neighbors(
          input_maps, candidate_maps, distances, num_matches=2
      )
      assert input_maps_filtered.shape == (3, 2, 2)
      assert input_maps_filtered.equal(
          torch.tensor(
              [
                  [[1, 4], [2, 5]],
                  [[8, 11], [7, 10]],
                  [[13, 16], [14, 17]],
              ]
          )
      )
      assert candidate_maps_filtered.shape == (3, 2, 2)
      assert candidate_maps_filtered.equal(
          torch.tensor(
              [[[1, 1], [2, 2]], [[3, 3], [3, 3]], [[1, 1], [1, 1]]]
          )
      )
      # Test num_matches = 1
      input_maps_filtered, candidate_maps_filtered = nearest_neighbors(
          input_maps, candidate_maps, distances, num_matches=1
      )
      assert input_maps_filtered.shape == (3, 1, 2)
      assert input_maps_filtered.equal(
          torch.tensor([[[1, 4]], [[8, 11]], [[13, 16]]])
      )
      assert candidate_maps_filtered.shape == (3, 1, 2)
      assert candidate_maps_filtered.equal(
          torch.tensor([[[1, 1]], [[3, 3]], [[1, 1]]])
      )
  
imports: import copy
  import unittest
  import torch
  import torch.nn
  from lightly.models import utils
  from lightly.models.utils import _no_grad_trunc_normal, activate_requires_grad, batch_shuffle, batch_unshuffle, deactivate_requires_grad, nearest_neighbors, normalize_weight, update_momentum
label_desc: Use image node_id 18 for calling the TestModelUtils obj's underlying member method code with example usage: obj.test_nearest_neighbors() without return types