Shen Li 06aaf8c20d Add set_device_map to TensorPipeOptions to support GPU args (#42637) ... Summary: Pull Request resolved: https://github.com/pytorch/pytorch/pull/42637 This commit enables sending non-CPU tensors through RPC using TensorPipe backend. Users can configure device mappings by calling set_map_location on `TensorPipeRpcBackendOptions`. Internally, the `init_rpc` API verifies the correctness of device mappings. It will shutdown RPC if the check failed, or proceed and pass global mappings to `TensorPipeAgent` if the check was successful. For serde, we added a device indices field to TensorPipe read and write buffers, which should be either empty (all tensors must be on CPU) or match the tensors in order and number in the RPC message. This commit does not yet avoid zero-copy, the tensor is always moved to CPU on the sender and then moved to the specified device on the receiver. Test Plan: Imported from OSS Reviewed By: izdeby Differential Revision: D23011572 Pulled By: mrshenli fbshipit-source-id: 62b617eed91237d4e9926bc8551db78b822a1187		2020-08-14 18:46:55 -07:00
..
_internal	Add set_device_map to TensorPipeOptions to support GPU args (#42637)	2020-08-14 18:46:55 -07:00
__init__.py	Adds list of operator-related information for testing (#41662)	2020-08-11 11:34:53 -07:00