Maybe manually implementing a specialized function would work well for you.
The implementation of checkpointing is a bit tricky, but grabbing CheckpointFunction and amending it to support your stats should not be too hard.
Best regards
Thomas