Shortcuts

DeviceDtypeModuleMixin

class pytorch_lightning.core.mixins.DeviceDtypeModuleMixin[source]

Bases: torch.nn.Module

cpu()[source]

Moves all model parameters and buffers to the CPU.

Returns

self

Return type

Module

cuda(device=None)[source]

Moves all model parameters and buffers to the GPU. This also makes associated parameters and buffers different objects. So it should be called before constructing optimizer if the module will live on GPU while being optimized.

Parameters

device (Union[device, int, None]) – if specified, all parameters will be copied to that device

Returns

self

Return type

Module

double()[source]

Casts all floating point parameters and buffers to double datatype.

Returns

self

Return type

Module

float()[source]

Casts all floating point parameters and buffers to float datatype.

Returns

self

Return type

Module

half()[source]

Casts all floating point parameters and buffers to half datatype.

Returns

self

Return type

Module

to(*args, **kwargs)[source]

Moves and/or casts the parameters and buffers.

This can be called as .. function:: to(device=None, dtype=None, non_blocking=False) .. function:: to(dtype, non_blocking=False) .. function:: to(tensor, non_blocking=False) Its signature is similar to torch.Tensor.to(), but only accepts floating point desired dtype s. 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. When non_blocking is set, it tries to convert/move asynchronously with respect to the host if possible, e.g., moving CPU Tensors with pinned memory to CUDA devices. See below for examples.

Note

This method modifies the module in-place.

Parameters
  • device – the desired device of the parameters and buffers in this module

  • dtype – the desired floating point type of the floating point parameters and buffers in this module

  • 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

Example::
>>> class ExampleModule(DeviceDtypeModuleMixin):
...     def __init__(self, weight: torch.Tensor):
...         super().__init__()
...         self.register_buffer('weight', weight)
>>> _ = torch.manual_seed(0)
>>> module = ExampleModule(torch.rand(3, 4))
>>> module.weight 
tensor([[...]])
>>> module.to(torch.double)
ExampleModule()
>>> module.weight 
tensor([[...]], dtype=torch.float64)
>>> cpu = torch.device('cpu')
>>> module.to(cpu, dtype=torch.half, non_blocking=True)
ExampleModule()
>>> module.weight 
tensor([[...]], dtype=torch.float16)
>>> module.to(cpu)
ExampleModule()
>>> module.weight 
tensor([[...]], dtype=torch.float16)
>>> module.device
device(type='cpu')
>>> module.dtype
torch.float16
type(dst_type)[source]

Casts all parameters and buffers to dst_type.

Parameters

dst_type (type or string) – the desired type

Returns

self

Return type

Module

Read the Docs v: stable
Versions
latest
stable
1.6.3
1.6.2
1.6.1
1.6.0
1.5.10
1.5.9
1.5.8
1.5.7
1.5.6
1.5.5
1.5.4
1.5.3
1.5.2
1.5.1
1.5.0
1.4.9
1.4.8
1.4.7
1.4.6
1.4.5
1.4.4
1.4.3
1.4.2
1.4.1
1.4.0
1.3.8
1.3.7
1.3.6
1.3.5
1.3.4
1.3.3
1.3.2
1.3.1
1.3.0
1.2.10
1.2.8
1.2.7
1.2.6
1.2.5
1.2.4
1.2.3
1.2.2
1.2.1
1.2.0
1.1.8
1.1.7
1.1.6
1.1.5
1.1.4
1.1.3
1.1.2
1.1.1
1.1.0
1.0.8
1.0.7
1.0.6
1.0.5
1.0.4
1.0.3
1.0.2
1.0.1
1.0.0
0.10.0
0.9.0
0.8.5
0.8.4
0.8.3
0.8.2
0.8.1
0.8.0
0.7.6
0.7.5
0.7.4
0.7.3
0.7.2
0.7.1
0.7.0
0.6.0
0.5.3
0.4.9
docs_2
Downloads
html
On Read the Docs
Project Home
Builds

Free document hosting provided by Read the Docs.