Shortcuts

decorators

Functions

auto_move_data

Decorator for LightningModule methods for which input arguments should be moved automatically to the correct 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

Read the Docs v: 1.0.7
Versions
latest
stable
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.2
0.5.3
0.4.9
release-1.0.x
Downloads
On Read the Docs
Project Home
Builds

Free document hosting provided by Read the Docs.