Skip to content

Algorithm

StepOutputDict #

Bases: TypedDict

A dictionary that shows what an Algorithm can output from training/validation/test_step.

loss instance-attribute #

loss: NotRequired[Tensor | float]

Optional loss tensor that can be returned by those methods.

Algorithm #

Bases: Module, Protocol[BatchType, StepOutputType]

Protocol that adds more type information to the lightning.LightningModule class.

This adds some type information on top of the LightningModule class, namely: - BatchType: The type of batch that is produced by the dataloaders of the datamodule - StepOutputType, the output type created by the step methods.

The networks themselves are created separately and passed as a constructor argument. This is meant to make it easier to compare different learning algorithms on the same network architecture.