ref: inner train loop (intermediate step) 3/n

Author: williamFalcon

pytorch_lightning/trainer/training_loop.py

@@ -938,8 +938,7 @@ def optimizer_closure(self, split_batch, batch_idx, opt_idx, optimizer, hiddens)
         result = self.train_loop.training_step(split_batch, batch_idx, opt_idx, hiddens)
 
         # backward pass
-        with self.profiler.profile('model_backward'):
-            result.closure_loss = self.accelerator_backend.backward(result.closure_loss, optimizer, opt_idx)
+        self.train_loop.backward(result, optimizer, opt_idx)
 
         # hook
         self.train_loop.on_after_backward(result.training_step_output, batch_idx, result.untouched_loss)

pytorch_lightning/trainer/training_loop_temp.py

@@ -134,6 +134,11 @@ def get_optimizers_iterable(self):
         opt_idx = np.argmax(optimizer_freq_cumsum > current_place_in_loop)
         return [(opt_idx, self.trainer.optimizers[opt_idx])]
 
+    def backward(self, result, optimizer, opt_idx):
+        # backward pass
+        with self.trainer.profiler.profile('model_backward'):
+            result.closure_loss = self.trainer.accelerator_backend.backward(result.closure_loss, optimizer, opt_idx)
+
     def on_after_backward(self, training_step_output, batch_idx, untouched_loss):
         is_result_obj = isinstance(training_step_output, Result)