TrainerInfo

class stable_pretraining.callbacks.TrainerInfo

Bases: Callback

Links the trainer to the DataModule for enhanced functionality.

This callback establishes a bidirectional connection between the trainer and DataModule, enabling the DataModule to access trainer information such as device placement, distributed training state, and other runtime configurations.

This is particularly useful for DataModules that need to adapt their behavior based on trainer configuration (e.g., device-aware data loading, distributed sampling adjustments).

Note

Only works with DataModule instances that have a set_pl_trainer method. A warning is logged if using a custom DataModule without this method.

setup(trainer, pl_module, stage)

Called when fit, validate, test, predict, or tune begins.