Allow user to specify 'step' key while logging metrics (#808)

* allow to specify 'step' key

* add test

* docs to log_metrics

* fix test

* rename

* also rename

tests/test_logging.py

@@ -376,3 +376,33 @@ def test_custom_logger(tmpdir):
     assert logger.hparams_logged == hparams
     assert logger.metrics_logged != {}
     assert logger.finalized_status == "success"
+
+
+def test_adding_step_key(tmpdir):
+    logged_step = 0
+
+    def _validation_end(outputs):
+        nonlocal logged_step
+        logged_step += 1
+        return {"log": {"step": logged_step, "val_acc": logged_step / 10}}
+
+    def _log_metrics_decorator(log_metrics_fn):
+        def decorated(metrics, step):
+            if "val_acc" in metrics:
+                assert step == logged_step
+            return log_metrics_fn(metrics, step)
+
+        return decorated
+
+    model, hparams = tutils.get_model()
+    model.validation_end = _validation_end
+    trainer_options = dict(
+        max_epochs=4,
+        default_save_path=tmpdir,
+        train_percent_check=0.001,
+        val_percent_check=0.01,
+        num_sanity_val_steps=0
+    )
+    trainer = Trainer(**trainer_options)
+    trainer.logger.log_metrics = _log_metrics_decorator(trainer.logger.log_metrics)
+    trainer.fit(model)