code: def test_knn() -> None:
      feature_bank = torch.tensor(
          [
              [1.0, 1.0, 1.0],
              [-1.0, 1.0, 1.0],
              [-1.0, -1.0, 1.0],
              [-1.0, -1.0, -1.0],
          ]
      ).t()
      feature_labels = torch.tensor([0, 1, 2, 3])
      features = torch.tensor(
          [
              [1.0, 1.0, 1.0],
              [-1.0, 1.1, -1.0],
          ]
      )
      feature_bank = F.normalize(feature_bank, dim=0)
      features = F.normalize(features, dim=1)
      pred_labels = knn.knn_predict(
          feature=features,
          feature_bank=feature_bank,
          feature_labels=feature_labels,
          num_classes=4,
          knn_k=4,
      )
      assert pred_labels.tolist() == [
          [0, 1, 2, 3],
          [1, 3, 0, 2],
      ]
  
imports: import torch
  import torch.nn.functional
  from lightly.utils.benchmarking import knn
label_desc: Use image node_id 1 for calling a global function with example usage: test_knn() without return types