diff --git a/pytorch_lightning/trainer/model_hooks_mixin.py b/pytorch_lightning/trainer/model_hooks_mixin.py index 23c82485..442c5209 100644 --- a/pytorch_lightning/trainer/model_hooks_mixin.py +++ b/pytorch_lightning/trainer/model_hooks_mixin.py @@ -1,3 +1,4 @@ +import inspect from pytorch_lightning.core.lightning import LightningModule @@ -15,3 +16,8 @@ class TrainerModelHooksMixin(object): # when code pointers are different, it was overriden is_overriden = getattr(model, f_name).__code__ is not getattr(super_object, f_name).__code__ return is_overriden + + def has_arg(self, f_name, arg_name): + model = self.get_model() + f_op = getattr(model, f_name, None) + return arg_name in inspect.signature(f_op).parameters diff --git a/pytorch_lightning/trainer/train_loop_mixin.py b/pytorch_lightning/trainer/train_loop_mixin.py index 0bff3d59..49cd3d70 100644 --- a/pytorch_lightning/trainer/train_loop_mixin.py +++ b/pytorch_lightning/trainer/train_loop_mixin.py @@ -149,6 +149,7 @@ When this flag is enabled each batch is split into sequences of size truncated_b """ +import inspect import numpy as np from pytorch_lightning.utilities.debugging import MisconfigurationException @@ -434,8 +435,15 @@ class TrainerTrainLoopMixin(object): # --------------- # enable not needing to add opt_idx to training_step args = [batch, batch_idx] + if len(self.optimizers) > 1: - args.append(opt_idx) + if self.has_arg('training_step', 'optimizer_idx'): + args.append(opt_idx) + else: + raise ValueError( + f'Your LightningModule defines {len(self.optimizers)} optimizers but ' + f'training_step is missing the "optimizer_idx" argument.' + ) # pass hiddens if using tbptt if self.truncated_bptt_steps is not None: