code: def get_sorted_predictions(
      batch_size: int,
      num_classes: int,
      distances: torch.Tensor,
      retrieval_one_hot: torch.Tensor,
      temperature: float,
  ) -> torch.Tensor:
      """
      From the closest neighbors compute the top predicted classes
      by blending the probabilities of each neighbor's class
      """
      assert distances.shape[0] == batch_size
      num_neighbors = distances.shape[1]
      assert retrieval_one_hot.shape[0] == batch_size * num_neighbors
      assert retrieval_one_hot.shape[1] == num_classes
  
      distances = distances.clone().div_(temperature).exp_()
      probs = torch.sum(
          torch.mul(
              retrieval_one_hot.view(
                  batch_size, num_neighbors, num_classes
              ),
              distances.view(batch_size, num_neighbors, 1),
          ),
          dim=1,
      )
      _, predictions = probs.sort(dim=1, descending=True)
  
      assert probs.shape == torch.Size([batch_size, num_classes])
      assert predictions.shape == torch.Size([batch_size, num_classes])
      return predictions
  
imports: import logging
  from typing import List, NamedTuple
  import numpy
  import torch
  from torch import nn
  from vissl.config import AttrDict
  from vissl.hooks import default_hook_generator
  from vissl.models.model_helpers import get_trunk_output_feature_names
  from vissl.utils.checkpoint import get_checkpoint_folder
  from vissl.utils.distributed_launcher import create_submitit_executor, launch_distributed
  from vissl.utils.env import set_env_vars
  from vissl.utils.extract_features_utils import ExtractedFeaturesLoader
  from vissl.utils.hydra_config import print_cfg
  from vissl.utils.io import save_file
  from vissl.utils.logger import setup_logging, shutdown_logging
label_desc: Use image node_id 3 for calling a global function with example usage: get_sorted_predictions(batch_size, num_classes, distances, retrieval_one_hot, temperature) and returns: predictions