decorators

Functions

parameter_validation

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

pytorch_lightning.core.decorators.parameter_validation(fn)

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

Return type

Callable

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.

See also

• XLA Documentation