lightning-trainer-utils 2025.5.26.15.18

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

PyTorch Lightning Trainer Utilities

Installation

pip install lightning-trainer-utils

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]

Trainer

Global Step

batch_step = num_samples / (batch_size * num_devices) trainer_global_step = num_samples / (batch_size * num_devices * grad_accumulation) SaveCheckpoint also use trainer_global_step.