def _init_rpc(self):
    pg_rank = dist.get_rank()
    rpc_rank = rpc.get_worker_info().id
    if (pg_rank != rpc_rank):
        raise ValueError(f'Default ProcessGroup and RPC ranks must be the same for ShardedTensor, found process group rank: {pg_rank} and RPC rank: {rpc_rank}')
    self._remote_shards = {}
    worker_infos = rpc._get_current_rpc_agent().get_worker_infos()
    rank_to_name = {}
    name_to_rank = {}
    for worker_info in worker_infos:
        rank_to_name[worker_info.id] = worker_info.name
        name_to_rank[worker_info.name] = worker_info.id
    all_tensor_ids = rpc.api._all_gather(self._sharded_tensor_id)
    futs = []
    rpc_rank = rpc.get_worker_info().id
    for rank in range(dist.get_world_size()):
        if (rank == dist.get_rank()):
            continue
        if (len(self.local_shards()) != 0):
            rrefs: List[rpc.RRef[Shard]] = [rpc.RRef(shard) for shard in self.local_shards()]
            fut = rpc.rpc_async(rank, _register_remote_shards, args=(all_tensor_ids[rank_to_name[rank]], rrefs, rpc_rank))
            futs.append(fut)
    torch.futures.wait_all(futs)
    rpc.api._all_gather(None)
def _init_rpc(self):
    pg_rank = dist.get_rank()
    rpc_rank = rpc.get_worker_info().id
    if (pg_rank != rpc_rank):
        raise ValueError(f'Default ProcessGroup and RPC ranks must be the same for ShardedTensor, found process group rank: {pg_rank} and RPC rank: {rpc_rank}')
    self._remote_shards = {}
    worker_infos = rpc._get_current_rpc_agent().get_worker_infos()
    rank_to_name = {}
    name_to_rank = {}
    for worker_info in worker_infos:
        rank_to_name[worker_info.id] = worker_info.name
        name_to_rank[worker_info.name] = worker_info.id
    all_tensor_ids = rpc.api._all_gather(self._sharded_tensor_id)
    futs = []
    rpc_rank = rpc.get_worker_info().id
    for rank in range(dist.get_world_size()):
        if (rank == dist.get_rank()):
            continue
        if (len(self.local_shards()) != 0):
            rrefs: List[rpc.RRef[Shard]] = [rpc.RRef(shard) for shard in self.local_shards()]
            fut = rpc.rpc_async(rank, _register_remote_shards, args=(all_tensor_ids[rank_to_name[rank]], rrefs, rpc_rank))
            futs.append(fut)
    torch.futures.wait_all(futs)
    rpc.api._all_gather(None)