pytorch_pfn_extras.to#
- pytorch_pfn_extras.to(module_or_tensor, device, *, options=None, runtime_class=None, config=None)#
A function to transfer the given object to the given device.
If PyTorch’s device type is given as the
device
argument, the behavior of this function is equivalent tomodule_or_tensor.to(module_or_tensor, device)
.Otherwise, this function uses the Runtime mechanism. This function looks for the Runtime for the device from the RuntimeRegistry and delegates the actual transfer operation to it.
See also the documentation of
ppe.runtime.BaseRuntime
for details.- Parameters:
module_or_tensor (torch.nn.Module or torch.Tensor) – An object to be transferred.
device (torch.device or str) – The device that the input object is transferred to.
options (dict, optional) – An options of dictionary type that is passed to
runtime_class.__init__
as an argument.runtime_class (Optional[Type[BaseRuntime]]) – A runtime class inherited from BaseRuntime class. If
None
, a runtime class is automatically selected based on thedevice
argument from the runtime registry.config (dict, optional) – DEPRECATED. Use options.
- Returns:
A torch.Tensor with the specified device.
- Return type:
ModuleOrTensor