Rooba
/
lightning
mirror of https://github.com/Lightning-AI/lightning.git
1
0
Fork 0
Code Issues Packages Projects Releases Wiki Activity

Added optimizer_idx to backward call (#733)

This commit is contained in:
Anand Krishnamoorthy 2020-01-25 08:03:07 +09:00 committed by William Falcon
parent a804755e6e
commit 946aef6216
2 changed files with 4 additions and 2 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

3
pytorch_lightning/core/hooks.py
View File

@ -124,12 +124,13 @@ class ModelHooks(torch.nn.Module):
"""
pass
def backward(self, use_amp, loss, optimizer):
def backward(self, use_amp, loss, optimizer, optimizer_idx):
"""Override backward with your own implementation if you need to
:param use_amp: Whether amp was requested or not
:param loss: Loss is already scaled by accumulated grads
:param optimizer: Current optimizer being used
:param optimizer_idx: Index of the current optimizer being used
:return:
Called to perform backward step.

3
pytorch_lightning/trainer/training_loop.py
View File

@ -490,13 +490,14 @@ class TrainerTrainLoopMixin(ABC):
# backward pass
model_ref = self.get_model()
model_ref.backward(self.use_amp, closure_loss, optimizer)
model_ref.backward(self.use_amp, closure_loss, optimizer, opt_idx)
# track metrics for callbacks
all_callback_metrics.append(callback_metrics)
# track progress bar metrics
self.add_tqdm_metrics(progress_bar_metrics)
self.add_tqdm_metrics(progress_bar_metrics)
all_log_metrics.append(log_metrics)
# insert after step hook