code: def sharded_state_dict(self):
      sharded_state_dict = self.state_dict(keep_vars=True)
      # conv
      sharded_state_dict[
          "conv.weight"
      ] = ShardedTensor.from_rank_offsets(
          "conv.weight",
          sharded_state_dict["conv.weight"],
          (
              1,
              parallel_state.get_tensor_model_parallel_rank(),
              parallel_state.get_tensor_model_parallel_world_size(),
          ),
      )
      # bias is non-sharded
      sharded_state_dict["conv.bias"] = ShardedTensor.from_rank_offsets(
          "conv.bias", sharded_state_dict["conv.bias"]
      )
  
      # proj
      sharded_state_dict[
          "proj.weight"
      ] = ShardedTensor.from_rank_offsets(
          "proj.weight",
          sharded_state_dict["proj.weight"],
          (0, Utils.rank, Utils.world_size),
      )
      sharded_state_dict["proj.bias"] = ShardedTensor.from_rank_offsets(
          "proj.bias",
          sharded_state_dict["proj.bias"],
          (0, Utils.rank, Utils.world_size),
      )
      return sharded_state_dict
  
file_path: megatron-lm/tests/unit_tests/dist_checkpointing/test_optimizer.py
imports: import numpy
  import torch
  from torch.optim import Adam
  from megatron.core import parallel_state
  from megatron.core.dist_checkpointing import ShardedTensor, save, load
  from megatron.core.dist_checkpointing.dict_utils import nested_values
  from megatron.core.dist_checkpointing.optimizer import get_param_id_to_sharded_param_map, optim_state_to_sharding_state
  from megatron.core.dist_checkpointing.utils import extract_sharded_tensors
  from tests.unit_tests.dist_checkpointing import TempNamedDir
  from tests.unit_tests.test_utilities import Utils
label_desc: Use image node_id 2 for calling the Model obj's underlying member method code with example usage: obj.sharded_state_dict() and returns: sharded_state_dict