code: def load_weight(weight_file: str,
                  weight_name: str,
                  weight_file_cache: Dict[str, Dict]) -> mx.nd.NDArray:
      
      logger.info('Loading input weight file: %s', weight_file)
      if weight_file.endswith(".npy"):
          return np.load(weight_file)
      elif weight_file.endswith(".npz"):
          if weight_file not in weight_file_cache:
              weight_file_cache[weight_file] = np.load(weight_file)
          return weight_file_cache[weight_file][weight_name]
      else:
          if weight_file not in weight_file_cache:
              weight_file_cache[weight_file] = mx.nd.load(weight_file)
          return weight_file_cache[weight_file]['arg:%s' % weight_name].asnumpy()
docs: Load wight fron a file or the cache if it was loaded before.
  
      :param weight_file: Weight file.
      :param weight_name: Weight name.
      :param weight_file_cache: Cache of loaded files.
      :return: Loaded weight.