Why not include a callback that stops the training when the Differential Privacy budget has been spent? See [20853](https://github.com/Lightning-AI/pytorch-lightning/pull/20853). cc @lantiga @borda