Validates that the module parameter lengths match after moving to the device.
Validates that the module parameter lengths match after moving to the device. It is useful when tying weights on TPU’s.
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_devicemodel hook which is called after the module has been moved to the device.
- Return type