Shortcuts

decorators

Functions

auto_move_data

Decorator for LightningModule methods for which input arguments should be moved automatically to the correct device.

parameter_validation

Validates that the module parameter lengths match after moving to the device.

Decorator for LightningModule methods.

pytorch_lightning.core.decorators.auto_move_data(fn)[source]

Decorator for LightningModule methods for which input arguments should be moved automatically to the correct device. It as no effect if applied to a method of an object that is not an instance of LightningModule and is typically applied to __call__ or forward.

Parameters

fn (Callable) – A LightningModule method for which the arguments should be moved to the device the parameters are on.

Example:

# directly in the source code
class LitModel(LightningModule):

    @auto_move_data
    def forward(self, x):
        return x

# or outside
LitModel.forward = auto_move_data(LitModel.forward)

model = LitModel()
model = model.to('cuda')
model(torch.zeros(1, 3))

# input gets moved to device
# tensor([[0., 0., 0.]], device='cuda:0')
Return type

Callable

pytorch_lightning.core.decorators.parameter_validation(fn)[source]

Validates that the module parameter lengths match after moving to the device. It is useful when tying weights on TPU’s.

Parameters

fn (Callable) – model_to_device method

Note

TPU’s require weights to be tied/shared after moving the module to the device. Failure to do this results in the initialization of new weights which are not tied. To overcome this issue, weights should be tied using the on_post_move_to_device model hook which is called after the module has been moved to the device.

Return type

Callable

Read the Docs v: stable
Versions
latest
stable
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
ipynb-update
docs-search
Downloads
html
On Read the Docs
Project Home
Builds

Free document hosting provided by Read the Docs.