MolecularDiffusion.modules.models.tabasco.callbacks.dataset_stats

Attributes

log

Classes

DatasetStatsCallback

Sync dataset statistics and effective batch size to the Lightning model.

Module Contents

class MolecularDiffusion.modules.models.tabasco.callbacks.dataset_stats.DatasetStatsCallback

Bases: lightning.Callback

Sync dataset statistics and effective batch size to the Lightning model.

The datamodule must expose get_dataset_stats() and the model must implement set_data_stats(stats); otherwise the callback logs a warning and continues without raising.

on_fit_start(trainer, pl_module)

Hook called once at the start of training.

MolecularDiffusion.modules.models.tabasco.callbacks.dataset_stats.log