Metadata-Version: 2.4
Name: lightning_trainer_utils
Version: 2025.5.16.10.56
Summary: A Python package for using PyTorch Lightning with custom callbacks and model wrappers.
Home-page: https://github.com/manavmahan/lightning-trainer-utils
Author: Manav Mahan Singh
Author-email: manav@genaec.ai
Requires-Python: >=3.11
Description-Content-Type: text/markdown
Dynamic: author
Dynamic: author-email
Dynamic: description
Dynamic: description-content-type
Dynamic: home-page
Dynamic: requires-python
Dynamic: summary

[![PyPI version](https://badge.fury.io/py/lightning-trainer-utils.svg)](https://pypi.org/project/lightning-trainer-utils/)
# PyTorch Lightning Trainer Utilities

## ML Model Assumptions

### forward
- The model wrapper uses the forward function as follows:
```python
    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]
