def _register_remote_shards(self, remote_shards: List[rpc.RRef[Shard]], rpc_rank: int):
    self._remote_shards[rpc_rank] = remote_shards
def _register_remote_shards(self, remote_shards: List[rpc.RRef[Shard]], rpc_rank: int):
    self._remote_shards[rpc_rank] = remote_shards