allow user to control optimizer step for every optimizer

* added custom hook for user defined optimizer step * refactored to allow multiple optimizers different training_step * refactored to allow multiple optimizers different training_step * refactored to allow multiple optimizers different training_step * refactored to allow multiple optimizers different training_step * refactored to allow multiple optimizers different training_step * pep8
2026-09-11 12:31:23 +08:00 · 2019-08-13 09:32:45 -04:00
parent 1c08882e6c
commit 905a2e5a12
2 changed files with 124 additions and 83 deletions
@@ -66,6 +66,20 @@ class LightningModule(GradInformation, ModelIO, ModelHooks):
        """
        raise NotImplementedError

+    def optimizer_step(self, epoch_nb, batch_nb, optimizer, optimizer_i):
+        """
+        Do something instead of the standard optimizer behavior
+        :param epoch_nb:
+        :param batch_nb:
+        :param optimizer:
+        :param optimizer_i:
+        :return:
+        """
+        optimizer.step()
+
+        # clear gradients
+        optimizer.zero_grad()
+
    @data_loader
    def tng_dataloader(self):
        """