quaterion.train.callbacks.cleanup_callback module

class CleanupCallback[source]

Bases: Callback

teardown(trainer: Trainer, pl_module: TrainableModel, stage: Optional[str] = None) → None[source]

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