scvi.model.base.PyroJitGuideWarmup.on_train_epoch_end

PyroJitGuideWarmup.on_train_epoch_end(trainer, pl_module, unused=None)

Called when the train epoch ends.

To access all batch outputs at the end of the epoch, either:

  1. Implement training_epoch_end in the LightningModule and access outputs via the module OR

  2. Cache data across train batch hooks inside the callback implementation to post-process in this hook.