added on_after_backward

pytorch_lightning/models/trainer.py

@@ -775,6 +775,11 @@ class Trainer(TrainerIO):
         else:
             loss.backward()
 
+        # insert after step hook
+        if self.__is_function_implemented('on_after_backward'):
+            model_ref = self.__get_model()
+            response = model_ref.on_after_backward()
+
         if self.print_nan_grads:
             model = self.__get_model()
             for param in model.parameters():

pytorch_lightning/root_module/hooks.py

@@ -36,3 +36,10 @@ class ModelHooks(torch.nn.Module):
         """
         pass
 
+    def on_after_backward(self):
+        """
+        Called after loss.backward() and before optimizers do anything
+        :return:
+        """
+        pass
+