code: def state_dict_to_device(
      state_dict: Dict[str, Union[torch.Tensor, ShardedTensor]],
      pg: ProcessGroup,
      device: torch.device,
  ) -> Dict[str, Union[torch.Tensor, ShardedTensor]]:
      """
      Moves a state_dict to a device with a process group. Can handle ShardedTensors in the state_dict.
  
      Args:
          state_dict (Dict[str, Union[torch.Tensor, ShardedTensor]]): state_dict to move
          pg (ProcessGroup): Process Group used for comms
          device (torch.device): device to put state_dict on
      """
      ret = {}
      all_keys = state_dict_all_gather_keys(state_dict, pg)
      for key in all_keys:
          if key in state_dict:
              tensor = state_dict[key]
              if isinstance(tensor, ShardedTensor):
                  copied_shards = [
                      Shard.from_tensor_and_offsets(
                          tensor=shard.tensor.to(device),
                          shard_offsets=shard.metadata.shard_offsets,
                          rank=dist.get_rank(pg),
                      )
                      for shard in tensor.local_shards()
                  ]
                  ret[key] = ShardedTensor._init_from_local_shards(
                      copied_shards,
                      tensor.metadata().size,
                      process_group=pg,
                  )
              elif isinstance(tensor, torch.Tensor):
                  ret[key] = tensor.to(device)
              else:
                  raise ValueError(
                      f"Unsupported tensor {key} type {type(tensor)}"
                  )
          else:
              # No state_dict entries for table-wise sharding,
              # but need to follow full-sync.
              ret[key] = ShardedTensor._init_from_local_shards(
                  [],
                  [],
                  process_group=pg,
              )
      return ret
  
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 3 for calling a global function with example usage: state_dict_to_device(state_dict, pg, device) and returns: ret