code: def create_module_dict(
      *,
      module,
      n_per_class_list,
      n_tries,
      nb_knn,
      train_features,
      train_labels,
  ):
      modules = {}
      mapping = create_class_indices_mapping(train_labels)
      for npc in n_per_class_list:
          if npc < 0:  # Only one try needed when using the full data
              full_module = module(
                  train_features=train_features,
                  train_labels=train_labels,
                  nb_knn=nb_knn,
              )
              modules["full"] = ModuleDictWithForward(
                  {"1": full_module}
              )
              continue
          all_tries = {}
          for t in range(n_tries):
              final_indices = filter_train(mapping, npc, seed=t)
              k_list = list(set(nb_knn + [npc]))
              k_list = sorted([el for el in k_list if el <= npc])
              all_tries[str(t)] = module(
                  train_features=train_features[final_indices],
                  train_labels=train_labels[final_indices],
                  nb_knn=k_list,
              )
          modules[f"{npc} per class"] = ModuleDictWithForward(all_tries)
  
      return ModuleDictWithForward(modules)
  
imports: import argparse
  from functools import partial
  import json
  import logging
  import os
  import sys
  from typing import List, Optional
  import torch
  from torch.nn.functional import one_hot, softmax
  import dinov2.distributed
  from dinov2.data import SamplerType, make_data_loader, make_dataset
  from dinov2.data.transforms import make_classification_eval_transform
  from dinov2.eval.metrics import AccuracyAveraging, build_topk_accuracy_metric
  from dinov2.eval.setup import get_args_parser
  from dinov2.eval.setup import setup_and_build_model
  from dinov2.eval.utils import ModelWithNormalize, evaluate, extract_features
label_desc: Use image node_id 2 for calling a global function with example usage: create_module_dict() and returns: ModuleDictWithForward