lightning-trainer-utils 2025.5.15.13.39

A Python package for using PyTorch Lightning with custom callbacks and model wrappers.

PyTorch Lightning Trainer Utilities

ML Model Assumptions

forward

• The model wrapper uses the forward function as follows:

    output = self.model(**x, **self.forward_kwargs)
    return ModelOuput(**output)

It expects batch as dict and returns a dict with keys [loss, report, output].

return

• ML model should return a dict with the following keys:
  • loss
  • report
  • output [optional]