We are currently solving the following issue within our repository. Here is the issue text:
--- BEGIN ISSUE ---
Problems in TensorPipeRpcBackendOptions device mapping documentation?
## 📚 Documentation

The new release of PyTorch 1.8 introduces CUDA-support in RPC.
I've referred to the RPC documentation, and the only reference for the CUDA-support I could find is under [`TensorPipeRpcBackendOptions`](https://pytorch.org/docs/1.8.0/rpc.html#torch.distributed.rpc.TensorPipeRpcBackendOptions) and [`set_device_map`](https://pytorch.org/docs/1.8.0/rpc.html#torch.distributed.rpc.TensorPipeRpcBackendOptions.set_device_map).
Seems like setting up CUDA-support is simply done by supplying a device mapping in the `TensorPipeRpcBackendOptions`, pretty cool.

However, I find the documentation for the `device_maps`/`device_map` to be unclear. It seems that `TensorPipeRpcBackendOptions`'s `device_maps` is a dictionary where the keys are worker names, but I'm not exactly sure what the structure of the dictionary's values should be like? Supposedly each value should be some sort of dictionary (as indicated by the parameter's type - `Dict[str, Dict]`), yet the example code provides a set:  `device_maps={"worker1": {0, 1}}`. I don't really understand how does this "map worker0's cuda:0 to worker1's cuda:1"?

Same for `set_device_map`'s `device_map`, the parameter's type also indicates it's a dictionary (`(Dict of python:int, str, or torch.device)`), but doesn't quite explain its structure. And again, the example code provides a set: `options.set_device_map("worker1", {1, 2})`.

It is also not explained how to define a GPU->CPU mapping (or vice versa).

Apart for this, there are 2 obvious errors in the example code provided in that documentation:

1. There is a missing comma in the following part:
```python
>>> rpc.init_rpc(
>>>     "worker0",
>>>     rank=0,
>>>     world_size=2  # <-- missing comma
>>>     backend=rpc.BackendType.TENSORPIPE,
>>>     rpc_backend_options=options
>>> )
```
2. I don't see how it is possible that those two `print`s will give different results. I'm guessing that the second line should read `print(rets[1])`?
```python
>>> print(rets[0])  # tensor([2., 2.], device='cuda:0')
>>> print(rets[0])  # tensor([2., 2.], device='cuda:1')
```

cc @pietern @mrshenli @pritamdamania87 @zhaojuanmao @satgera @gqchen @aazzolini @rohan-varma @jjlilley @osalpekar @jiayisuse @mrzzd @agolynski @SciPioneer @H-Huang @cbalioglu

--- END ISSUE ---

Below are some code segments, each from a relevant file. One or more of these files may contain bugs.
--- BEGIN FILES ---
Path: `torch/distributed/rpc/options.py`
Content:
```
1 from torch._C._distributed_rpc import _TensorPipeRpcBackendOptionsBase
2 from . import constants as rpc_contants
3 
4 import torch
5 
6 from typing import Dict, List
7 
8 
9 class TensorPipeRpcBackendOptions(_TensorPipeRpcBackendOptionsBase):
10     r"""
11     The backend options for
12     :class:`~torch.distributed.rpc.TensorPipeAgent`, derived from
13     :class:`~torch.distributed.rpc.RpcBackendOptions`.
14 
15     Args:
16         num_worker_threads (int, optional): The number of threads in the
17             thread-pool used by
18             :class:`~torch.distributed.rpc.TensorPipeAgent` to execute
19             requests (default: 16).
20         rpc_timeout (float, optional): The default timeout, in seconds,
21             for RPC requests (default: 60 seconds). If the RPC has not
22             completed in this timeframe, an exception indicating so will
23             be raised. Callers can override this timeout for individual
24             RPCs in :meth:`~torch.distributed.rpc.rpc_sync` and
25             :meth:`~torch.distributed.rpc.rpc_async` if necessary.
26         init_method (str, optional): The URL to initialize the distributed
27             store used for rendezvous. It takes any value accepted for the
28             same argument of :meth:`~torch.distributed.init_process_group`
29             (default: ``env://``).
30         device_maps (Dict[str, Dict]): Device placement mappings from this
31             worker to the callee. Key is the callee worker name and value the
32             dictionary (``Dict`` of ``int``, ``str``, or ``torch.device``) that
33             maps this worker's devices to the callee worker's devices.
34             (default: ``None``)
35     """
36     def __init__(
37         self,
38         *,
39         num_worker_threads: int = rpc_contants.DEFAULT_NUM_WORKER_THREADS,
40         rpc_timeout: float = rpc_contants.DEFAULT_RPC_TIMEOUT_SEC,
41         init_method: str = rpc_contants.DEFAULT_INIT_METHOD,
42         device_maps: Dict = None,
43         _transports: List = None,
44         _channels: List = None,
45     ):
46         super().__init__(
47             num_worker_threads,
48             _transports,
49             _channels,
50             rpc_timeout,
51             init_method,
52             device_maps if device_maps else {}
53         )
54 
55     def set_device_map(self, to: str, device_map: Dict):
56         r"""
57         Set device mapping between each RPC caller and callee pair. This
58         function can be called multiple times to incrementally add
59         device placement configurations.
60 
61         Args:
62             worker_name (str): Callee name.
63             device_map (Dict of int, str, or torch.device): Device placement
64                 mappings from this worker to the callee. This map must be
65                 invertible.
66 
67         Example::
68             >>> # both workers
69             >>> def add(x, y):
70             >>>     print(x)  # tensor([1., 1.], device='cuda:1')
71             >>>     return x + y, (x + y).to(2)
72             >>>
73             >>> # on worker 0
74             >>> options = TensorPipeRpcBackendOptions(
75             >>>     num_worker_threads=8,
76             >>>     device_maps={"worker1": {0, 1}}
77             >>>     # maps worker0's cuda:0 to worker1's cuda:1
78             >>> )
79             >>> options.set_device_map("worker1", {1, 2})
80             >>> # maps worker0's cuda:1 to worker1's cuda:2
81             >>>
82             >>> rpc.init_rpc(
83             >>>     "worker0",
84             >>>     rank=0,
85             >>>     world_size=2
86             >>>     backend=rpc.BackendType.TENSORPIPE,
87             >>>     rpc_backend_options=options
88             >>> )
89             >>>
90             >>> x = torch.ones(2)
91             >>> rets = rpc.rpc_sync("worker1", add, args=(x.to(0), 1))
92             >>> # The first argument will be moved to cuda:1 on worker1. When
93             >>> # sending the return value back, it will follow the invert of
94             >>> # the device map, and hence will be moved back to cuda:0 and
95             >>> # cuda:1 on worker0
96             >>> print(rets[0])  # tensor([2., 2.], device='cuda:0')
97             >>> print(rets[0])  # tensor([2., 2.], device='cuda:1')
98         """
99         device_index_map = {}
100         curr_device_maps = super().device_maps
101         for k in device_map:
102             v = device_map[k]
103             k, v = torch.device(k), torch.device(v)
104             if k.type != 'cuda' or v.type != 'cuda':
105                 raise ValueError(
106                     "`set_device_map` only supports CUDA devices, "
107                     f"but got device pair {k}: {v}"
108 
109                 )
110             if to in curr_device_maps and k.index in curr_device_maps[to]:
111                 curr_v = super().device_maps[to][k.index]
112                 if curr_v != v.index:
113                     raise ValueError(
114                         "`set_device_map` only supports 1-to-1 mapping, "
115                         f"trying to map {k} to {v} and {curr_v}"
116                     )
117             device_index_map[k.index] = v.index
118         super().set_device_map(to, device_index_map)
119
```
--- END FILES ---

Please first localize the bug based on the issue statement, and then generate a patch according to the `git diff` format fenced by three backticks.

Here is an example:

```diff
diff --git a/examples/server_async.py b/examples/server_async.py
--- a/examples/server_async.py
+++ b/examples/server_async.py
@@ -313,4 +313,4 @@
 
 
 if __name__ == "__main__":
-    asyncio.run(run_async_server("."), debug=True)
+    asyncio.run(run_async_server(), debug=True)
diff --git a/examples/server_sync.py b/examples/server_sync.py
--- a/examples/server_sync.py
+++ b/examples/server_sync.py
@@ -313,5 +313,5 @@
 
 
 if __name__ == "__main__":
-    server = run_sync_server(".")
+    server = run_sync_server()
     server.shutdown()

```

diff --git a/torch/distributed/rpc/options.py b/torch/distributed/rpc/options.py
--- a/torch/distributed/rpc/options.py
+++ b/torch/distributed/rpc/options.py
@@ -73,16 +73,16 @@
             >>> # on worker 0
             >>> options = TensorPipeRpcBackendOptions(
             >>>     num_worker_threads=8,
-            >>>     device_maps={"worker1": {0, 1}}
+            >>>     device_maps={"worker1": {0: 1}}
             >>>     # maps worker0's cuda:0 to worker1's cuda:1
             >>> )
-            >>> options.set_device_map("worker1", {1, 2})
+            >>> options.set_device_map("worker1", {1: 2})
             >>> # maps worker0's cuda:1 to worker1's cuda:2
             >>>
             >>> rpc.init_rpc(
             >>>     "worker0",
             >>>     rank=0,
-            >>>     world_size=2
+            >>>     world_size=2,
             >>>     backend=rpc.BackendType.TENSORPIPE,
             >>>     rpc_backend_options=options
             >>> )
@@ -94,7 +94,7 @@
             >>> # the device map, and hence will be moved back to cuda:0 and
             >>> # cuda:1 on worker0
             >>> print(rets[0])  # tensor([2., 2.], device='cuda:0')
-            >>> print(rets[0])  # tensor([2., 2.], device='cuda:1')
+            >>> print(rets[1])  # tensor([2., 2.], device='cuda:1')
         """
         device_index_map = {}
         curr_device_maps = super().device_maps

{"golden_diff": "diff --git a/torch/distributed/rpc/options.py b/torch/distributed/rpc/options.py\n--- a/torch/distributed/rpc/options.py\n+++ b/torch/distributed/rpc/options.py\n@@ -73,16 +73,16 @@\n             >>> # on worker 0\n             >>> options = TensorPipeRpcBackendOptions(\n             >>>     num_worker_threads=8,\n-            >>>     device_maps={\"worker1\": {0, 1}}\n+            >>>     device_maps={\"worker1\": {0: 1}}\n             >>>     # maps worker0's cuda:0 to worker1's cuda:1\n             >>> )\n-            >>> options.set_device_map(\"worker1\", {1, 2})\n+            >>> options.set_device_map(\"worker1\", {1: 2})\n             >>> # maps worker0's cuda:1 to worker1's cuda:2\n             >>>\n             >>> rpc.init_rpc(\n             >>>     \"worker0\",\n             >>>     rank=0,\n-            >>>     world_size=2\n+            >>>     world_size=2,\n             >>>     backend=rpc.BackendType.TENSORPIPE,\n             >>>     rpc_backend_options=options\n             >>> )\n@@ -94,7 +94,7 @@\n             >>> # the device map, and hence will be moved back to cuda:0 and\n             >>> # cuda:1 on worker0\n             >>> print(rets[0])  # tensor([2., 2.], device='cuda:0')\n-            >>> print(rets[0])  # tensor([2., 2.], device='cuda:1')\n+            >>> print(rets[1])  # tensor([2., 2.], device='cuda:1')\n         \"\"\"\n         device_index_map = {}\n         curr_device_maps = super().device_maps\n", "issue": "Problems in TensorPipeRpcBackendOptions device mapping documentation?\n## \ud83d\udcda Documentation\r\n\r\nThe new release of PyTorch 1.8 introduces CUDA-support in RPC.\r\nI've referred to the RPC documentation, and the only reference for the CUDA-support I could find is under [`TensorPipeRpcBackendOptions`](https://pytorch.org/docs/1.8.0/rpc.html#torch.distributed.rpc.TensorPipeRpcBackendOptions) and [`set_device_map`](https://pytorch.org/docs/1.8.0/rpc.html#torch.distributed.rpc.TensorPipeRpcBackendOptions.set_device_map).\r\nSeems like setting up CUDA-support is simply done by supplying a device mapping in the `TensorPipeRpcBackendOptions`, pretty cool.\r\n\r\nHowever, I find the documentation for the `device_maps`/`device_map` to be unclear. It seems that `TensorPipeRpcBackendOptions`'s `device_maps` is a dictionary where the keys are worker names, but I'm not exactly sure what the structure of the dictionary's values should be like? Supposedly each value should be some sort of dictionary (as indicated by the parameter's type - `Dict[str, Dict]`), yet the example code provides a set:  `device_maps={\"worker1\": {0, 1}}`. I don't really understand how does this \"map worker0's cuda:0 to worker1's cuda:1\"?\r\n\r\nSame for `set_device_map`'s `device_map`, the parameter's type also indicates it's a dictionary (`(Dict of python:int, str, or torch.device)`), but doesn't quite explain its structure. And again, the example code provides a set: `options.set_device_map(\"worker1\", {1, 2})`.\r\n\r\nIt is also not explained how to define a GPU->CPU mapping (or vice versa).\r\n\r\nApart for this, there are 2 obvious errors in the example code provided in that documentation:\r\n\r\n1. There is a missing comma in the following part:\r\n```python\r\n>>> rpc.init_rpc(\r\n>>>     \"worker0\",\r\n>>>     rank=0,\r\n>>>     world_size=2  # <-- missing comma\r\n>>>     backend=rpc.BackendType.TENSORPIPE,\r\n>>>     rpc_backend_options=options\r\n>>> )\r\n```\r\n2. I don't see how it is possible that those two `print`s will give different results. I'm guessing that the second line should read `print(rets[1])`?\r\n```python\r\n>>> print(rets[0])  # tensor([2., 2.], device='cuda:0')\r\n>>> print(rets[0])  # tensor([2., 2.], device='cuda:1')\r\n```\n\ncc @pietern @mrshenli @pritamdamania87 @zhaojuanmao @satgera @gqchen @aazzolini @rohan-varma @jjlilley @osalpekar @jiayisuse @mrzzd @agolynski @SciPioneer @H-Huang @cbalioglu\n", "before_files": [{"content": "from torch._C._distributed_rpc import _TensorPipeRpcBackendOptionsBase\nfrom . import constants as rpc_contants\n\nimport torch\n\nfrom typing import Dict, List\n\n\nclass TensorPipeRpcBackendOptions(_TensorPipeRpcBackendOptionsBase):\n    r\"\"\"\n    The backend options for\n    :class:`~torch.distributed.rpc.TensorPipeAgent`, derived from\n    :class:`~torch.distributed.rpc.RpcBackendOptions`.\n\n    Args:\n        num_worker_threads (int, optional): The number of threads in the\n            thread-pool used by\n            :class:`~torch.distributed.rpc.TensorPipeAgent` to execute\n            requests (default: 16).\n        rpc_timeout (float, optional): The default timeout, in seconds,\n            for RPC requests (default: 60 seconds). If the RPC has not\n            completed in this timeframe, an exception indicating so will\n            be raised. Callers can override this timeout for individual\n            RPCs in :meth:`~torch.distributed.rpc.rpc_sync` and\n            :meth:`~torch.distributed.rpc.rpc_async` if necessary.\n        init_method (str, optional): The URL to initialize the distributed\n            store used for rendezvous. It takes any value accepted for the\n            same argument of :meth:`~torch.distributed.init_process_group`\n            (default: ``env://``).\n        device_maps (Dict[str, Dict]): Device placement mappings from this\n            worker to the callee. Key is the callee worker name and value the\n            dictionary (``Dict`` of ``int``, ``str``, or ``torch.device``) that\n            maps this worker's devices to the callee worker's devices.\n            (default: ``None``)\n    \"\"\"\n    def __init__(\n        self,\n        *,\n        num_worker_threads: int = rpc_contants.DEFAULT_NUM_WORKER_THREADS,\n        rpc_timeout: float = rpc_contants.DEFAULT_RPC_TIMEOUT_SEC,\n        init_method: str = rpc_contants.DEFAULT_INIT_METHOD,\n        device_maps: Dict = None,\n        _transports: List = None,\n        _channels: List = None,\n    ):\n        super().__init__(\n            num_worker_threads,\n            _transports,\n            _channels,\n            rpc_timeout,\n            init_method,\n            device_maps if device_maps else {}\n        )\n\n    def set_device_map(self, to: str, device_map: Dict):\n        r\"\"\"\n        Set device mapping between each RPC caller and callee pair. This\n        function can be called multiple times to incrementally add\n        device placement configurations.\n\n        Args:\n            worker_name (str): Callee name.\n            device_map (Dict of int, str, or torch.device): Device placement\n                mappings from this worker to the callee. This map must be\n                invertible.\n\n        Example::\n            >>> # both workers\n            >>> def add(x, y):\n            >>>     print(x)  # tensor([1., 1.], device='cuda:1')\n            >>>     return x + y, (x + y).to(2)\n            >>>\n            >>> # on worker 0\n            >>> options = TensorPipeRpcBackendOptions(\n            >>>     num_worker_threads=8,\n            >>>     device_maps={\"worker1\": {0, 1}}\n            >>>     # maps worker0's cuda:0 to worker1's cuda:1\n            >>> )\n            >>> options.set_device_map(\"worker1\", {1, 2})\n            >>> # maps worker0's cuda:1 to worker1's cuda:2\n            >>>\n            >>> rpc.init_rpc(\n            >>>     \"worker0\",\n            >>>     rank=0,\n            >>>     world_size=2\n            >>>     backend=rpc.BackendType.TENSORPIPE,\n            >>>     rpc_backend_options=options\n            >>> )\n            >>>\n            >>> x = torch.ones(2)\n            >>> rets = rpc.rpc_sync(\"worker1\", add, args=(x.to(0), 1))\n            >>> # The first argument will be moved to cuda:1 on worker1. When\n            >>> # sending the return value back, it will follow the invert of\n            >>> # the device map, and hence will be moved back to cuda:0 and\n            >>> # cuda:1 on worker0\n            >>> print(rets[0])  # tensor([2., 2.], device='cuda:0')\n            >>> print(rets[0])  # tensor([2., 2.], device='cuda:1')\n        \"\"\"\n        device_index_map = {}\n        curr_device_maps = super().device_maps\n        for k in device_map:\n            v = device_map[k]\n            k, v = torch.device(k), torch.device(v)\n            if k.type != 'cuda' or v.type != 'cuda':\n                raise ValueError(\n                    \"`set_device_map` only supports CUDA devices, \"\n                    f\"but got device pair {k}: {v}\"\n\n                )\n            if to in curr_device_maps and k.index in curr_device_maps[to]:\n                curr_v = super().device_maps[to][k.index]\n                if curr_v != v.index:\n                    raise ValueError(\n                        \"`set_device_map` only supports 1-to-1 mapping, \"\n                        f\"trying to map {k} to {v} and {curr_v}\"\n                    )\n            device_index_map[k.index] = v.index\n        super().set_device_map(to, device_index_map)\n", "path": "torch/distributed/rpc/options.py"}], "after_files": [{"content": "from torch._C._distributed_rpc import _TensorPipeRpcBackendOptionsBase\nfrom . import constants as rpc_contants\n\nimport torch\n\nfrom typing import Dict, List\n\n\nclass TensorPipeRpcBackendOptions(_TensorPipeRpcBackendOptionsBase):\n    r\"\"\"\n    The backend options for\n    :class:`~torch.distributed.rpc.TensorPipeAgent`, derived from\n    :class:`~torch.distributed.rpc.RpcBackendOptions`.\n\n    Args:\n        num_worker_threads (int, optional): The number of threads in the\n            thread-pool used by\n            :class:`~torch.distributed.rpc.TensorPipeAgent` to execute\n            requests (default: 16).\n        rpc_timeout (float, optional): The default timeout, in seconds,\n            for RPC requests (default: 60 seconds). If the RPC has not\n            completed in this timeframe, an exception indicating so will\n            be raised. Callers can override this timeout for individual\n            RPCs in :meth:`~torch.distributed.rpc.rpc_sync` and\n            :meth:`~torch.distributed.rpc.rpc_async` if necessary.\n        init_method (str, optional): The URL to initialize the distributed\n            store used for rendezvous. It takes any value accepted for the\n            same argument of :meth:`~torch.distributed.init_process_group`\n            (default: ``env://``).\n        device_maps (Dict[str, Dict]): Device placement mappings from this\n            worker to the callee. Key is the callee worker name and value the\n            dictionary (``Dict`` of ``int``, ``str``, or ``torch.device``) that\n            maps this worker's devices to the callee worker's devices.\n            (default: ``None``)\n    \"\"\"\n    def __init__(\n        self,\n        *,\n        num_worker_threads: int = rpc_contants.DEFAULT_NUM_WORKER_THREADS,\n        rpc_timeout: float = rpc_contants.DEFAULT_RPC_TIMEOUT_SEC,\n        init_method: str = rpc_contants.DEFAULT_INIT_METHOD,\n        device_maps: Dict = None,\n        _transports: List = None,\n        _channels: List = None,\n    ):\n        super().__init__(\n            num_worker_threads,\n            _transports,\n            _channels,\n            rpc_timeout,\n            init_method,\n            device_maps if device_maps else {}\n        )\n\n    def set_device_map(self, to: str, device_map: Dict):\n        r\"\"\"\n        Set device mapping between each RPC caller and callee pair. This\n        function can be called multiple times to incrementally add\n        device placement configurations.\n\n        Args:\n            worker_name (str): Callee name.\n            device_map (Dict of int, str, or torch.device): Device placement\n                mappings from this worker to the callee. This map must be\n                invertible.\n\n        Example::\n            >>> # both workers\n            >>> def add(x, y):\n            >>>     print(x)  # tensor([1., 1.], device='cuda:1')\n            >>>     return x + y, (x + y).to(2)\n            >>>\n            >>> # on worker 0\n            >>> options = TensorPipeRpcBackendOptions(\n            >>>     num_worker_threads=8,\n            >>>     device_maps={\"worker1\": {0: 1}}\n            >>>     # maps worker0's cuda:0 to worker1's cuda:1\n            >>> )\n            >>> options.set_device_map(\"worker1\", {1: 2})\n            >>> # maps worker0's cuda:1 to worker1's cuda:2\n            >>>\n            >>> rpc.init_rpc(\n            >>>     \"worker0\",\n            >>>     rank=0,\n            >>>     world_size=2,\n            >>>     backend=rpc.BackendType.TENSORPIPE,\n            >>>     rpc_backend_options=options\n            >>> )\n            >>>\n            >>> x = torch.ones(2)\n            >>> rets = rpc.rpc_sync(\"worker1\", add, args=(x.to(0), 1))\n            >>> # The first argument will be moved to cuda:1 on worker1. When\n            >>> # sending the return value back, it will follow the invert of\n            >>> # the device map, and hence will be moved back to cuda:0 and\n            >>> # cuda:1 on worker0\n            >>> print(rets[0])  # tensor([2., 2.], device='cuda:0')\n            >>> print(rets[1])  # tensor([2., 2.], device='cuda:1')\n        \"\"\"\n        device_index_map = {}\n        curr_device_maps = super().device_maps\n        for k in device_map:\n            v = device_map[k]\n            k, v = torch.device(k), torch.device(v)\n            if k.type != 'cuda' or v.type != 'cuda':\n                raise ValueError(\n                    \"`set_device_map` only supports CUDA devices, \"\n                    f\"but got device pair {k}: {v}\"\n\n                )\n            if to in curr_device_maps and k.index in curr_device_maps[to]:\n                curr_v = super().device_maps[to][k.index]\n                if curr_v != v.index:\n                    raise ValueError(\n                        \"`set_device_map` only supports 1-to-1 mapping, \"\n                        f\"trying to map {k} to {v} and {curr_v}\"\n                    )\n            device_index_map[k.index] = v.index\n        super().set_device_map(to, device_index_map)\n", "path": "torch/distributed/rpc/options.py"}]}