code: def load(config, prefix: str, weights):
      if weights.process_group.size() > 1:
          try:
              weight = weights.get_sharded(f"{prefix}.weight", dim=0)
              should_gather = True
          except AssertionError:
              # If the vocab size is not divisible by number of shards
              # just load the entire thing.
              weight = weights.get_tensor(f"{prefix}.weight")
              should_gather = False
      else:
          weight = weights.get_tensor(f"{prefix}.weight")
          should_gather = False
  
      # GPTQ,AWQ,EETQ don't quantize heads (nor embeddings)
      if config.quantize in ["gptq", "awq", "eetq"]:
          quantize = None
      else:
          quantize = config.quantize
      return TensorParallelHead(
          get_linear(weight, bias=None, quantize=quantize),
          process_group=weights.process_group,
          should_gather=should_gather,
      )
  
imports: import os
  import torch
  import torch.distributed
  from torch import nn
  from torch.nn import functional
  from typing import List
  from loguru import logger
  from functools import lru_cache
  from accelerate import init_empty_weights
  from text_generation_server.utils.gptq.quant_linear import QuantLinear
  from text_generation_server.utils.import_utils import IS_CUDA_SYSTEM, IS_ROCM_SYSTEM
  from text_generation_server.utils.log import log_once
label_desc: Use image node_id 2 for calling the TensorParallelHead obj's underlying member method code with example usage: obj.load(config, prefix, weights) and returns: TensorParallelHead