oneflow.nn.Module.to

Module.to(device: Optional[Union[oneflow._oneflow_internal.device, str, int]] = ..., dtype: Optional[oneflow._oneflow_internal.dtype] = ...) → T
Module.to(dtype: oneflow._oneflow_internal.dtype) → T
Module.to(tensor: oneflow.Tensor) → T

Moves and/or casts the parameters and buffers.

This can be called as

to(device=None, dtype=None)
oneflow.nn.to(dtype)
oneflow.nn.to(memory_format=None)
oneflow.nn.to(tensor)

Its signature is similar to oneflow.Tensor.to(), but only accepts floating point dtypes. In addition, this method will only cast the floating point parameters and buffers to dtype (if given). The integral parameters and buffers will be moved device, if that is given, but with dtypes unchanged.

See below for examples.

Note

This method modifies the module in-place.

Parameters

  • device (oneflow.device) – the desired device of the parameters and buffers in this module

  • dtype (oneflow.dtype) – the desired floating point dtype of the parameters and buffers in this module

  • memory_format (oneflow.memory_format) – the desired memory format for 4D parameters and buffers in this module (keyword only argument)

  • tensor (oneflow.Tensor) – Tensor whose dtype and device are the desired dtype and device for all parameters and buffers in this module

Returns

self

Return type

Module

Examples:

>>> import oneflow as flow
>>> import oneflow.nn as nn
>>> linear = nn.Linear(2, 2)
>>> linear.weight.device
device(type='cpu', index=0)
>>> linear.weight.dtype
oneflow.float32
>>> linear.to(flow.double)
Linear(in_features=2, out_features=2, bias=True)
>>> linear.weight.dtype
oneflow.float64
>>> gpu1 = flow.device("cuda:1")
>>> linear.to(gpu1, dtype=flow.half)
Linear(in_features=2, out_features=2, bias=True)
>>> linear.weight.device
device(type='cuda', index=1)
>>> linear.weight.dtype
oneflow.float16
>>> cpu = flow.device("cpu")
>>> linear.to(cpu)
Linear(in_features=2, out_features=2, bias=True)
>>> linear.weight.device
device(type='cpu', index=0)