code: def get_tensor_shard(self, var, dim):
      world_size = self.process_group.size()
      rank = self.process_group.rank()
      block_size = var.size()[dim] // world_size
      start = rank * block_size
      stop = (rank + 1) * block_size
      if dim == 0:
          tensor = var[start:stop]
      elif dim == 1:
          tensor = var[:, start:stop]
      else:
          raise NotImplementedError(
              "Let's make that generic when needed"
          )
      tensor = tensor.to(dtype=self.dtype)
      tensor = tensor.to(device=self.device)
      return tensor
  
imports: import os
  from pathlib import Path
  from typing import List, Dict, Optional, Tuple
  from safetensors import safe_open, SafetensorError
  import torch
  from loguru import logger
  from huggingface_hub import hf_hub_download
  import json
  from text_generation_server.utils.log import log_once
label_desc: Use image node_id 12 for calling the Weights obj's underlying member method code with example usage: obj.get_tensor_shard(var, dim) and returns: tensor