inspect training_step for opt_idx (#573)

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

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: