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 pointdtype
s. In addition, this method will only cast the floating point parameters and buffers todtype
(if given). The integral parameters and buffers will be moveddevice
, 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 moduledtype (
oneflow.dtype
) – the desired floating point dtype of the parameters and buffers in this modulememory_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
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)
-