Add on_checkpoint decorator (#547)

* Add on_checkpoint decorator

* Update changelog

* Update docstrings

CHANGELOG.md

@@ -15,6 +15,7 @@ The format is based on [Keep a Changelog](http://keepachangelog.com/en/1.0.0/).
 - Added the option to give the ``only_if`` callback decorator a function of self and state rather than just state
 - Added Layer-sequential unit-variance (LSUV) initialization
 - Added ClassAppearanceModel callback and example page for visualising CNNs
+- Added on_checkpoint callback decorator
 ### Changed
 - `No_grad` and `enable_grad` decorators are now also context managers
 ### Deprecated

tests/callbacks/test_decorators.py

@@ -123,6 +123,12 @@ def example(state):
         state = 'test'
         self.assertTrue(callbacks.on_step_validation(example).on_step_validation(state) == state)
 
+    def test_on_checkpoint(self):
+        def example(state):
+            return state
+        state = 'test'
+        self.assertTrue(callbacks.on_checkpoint(example).on_checkpoint(state) == state)
+
     def test_add_to_loss(self):
         def example(state):
             return 1

torchbearer/callbacks/decorators.py

@@ -256,6 +256,19 @@ def on_step_validation(func):
     return bind_to(Callback.on_step_validation)(func)
 
 
+def on_checkpoint(func):
+    """ The :func:`on_checkpoint` decorator is used to initialise a :class:`.Callback` with :meth:`~.Callback.on_checkpoint`
+    calling the decorated function
+
+    Args:
+        func (function): The function(state) to *decorate*
+
+    Returns:
+        Callback: Initialised callback with :meth:`~.Callback.on_checkpoint` calling func
+    """
+    return bind_to(Callback.on_checkpoint)(func)
+
+
 def add_to_loss(func):
     """ The :func:`add_to_loss` decorator is used to initialise a :class:`.Callback` with the value returned from func
     being added to the loss