tabensemb.model.AbstractNN._register_sharded_tensor_state_dict_hooks_if_available

method

AbstractNN._register_sharded_tensor_state_dict_hooks_if_available() → None

Adds ShardedTensor state dict hooks if ShardedTensors are supported.

These hooks ensure that ShardedTensors are included when saving, and are loaded the LightningModule correctly.