code: def state_dict_gather(
      src: Dict[str, Union[torch.Tensor, ShardedTensor]],
      dst: Dict[str, torch.Tensor],
  ) -> None:
      """
      Gathers the values of the src state_dict of the keys present in the dst state_dict. Can handle ShardedTensors in the src state_dict.
  
      Args:
          src (Dict[str, Union[torch.Tensor, ShardedTensor]]): source's state_dict for this rank
          dst (Dict[str, torch.Tensor]): destination's state_dict
      """
      for key, dst_tensor in dst.items():
          src_tensor = src[key]
          if isinstance(src_tensor, ShardedTensor):
              src_tensor.gather(
                  out=dst_tensor if (dist.get_rank() == 0) else None
              )
          elif isinstance(src_tensor, torch.Tensor):
              dst_tensor.copy_(src_tensor)
          else:
              raise ValueError(
                  f"Unsupported tensor {key} type {type(src_tensor)}"
              )
  
imports: from typing import Dict, List, Union
  import torch
  from torch import distributed
  from torch.distributed import ProcessGroup
  from torch.distributed._shard.sharded_tensor import Shard, ShardedTensor
label_desc: Use image node_id 1 for calling a global function with example usage: state_dict_gather(src, dst) without return types