jactorch.cuda.copy#

Functions

async_copy_to(obj, dev[, main_stream])

Copy an object to a specific device asynchronizedly.

Functions

async_copy_to(obj, dev, main_stream=None)[source]#

Copy an object to a specific device asynchronizedly. If the param main_stream is provided, the copy stream will be synchronized with the main one.

Parameters:

obj (Iterable[Tensor] or Tensor) – a structure (e.g., a list or a dict) containing pytorch tensors.
dev (int) – the target device.
main_stream (stream) – the main stream to be synchronized.

Returns:

a deep copy of the data structure, with each tensor copied to the device.