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 to module_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 the device argument from the runtime registry.

  • config (dict, optional) – DEPRECATED. Use options.

Returns:

A torch.Tensor with the specified device.

Return type:

ModuleOrTensor