Avoid calling scheduler.step() after the last training iteration is done

Summary:
Pull Request resolved: https://github.com/facebookresearch/d2go/pull/358

Avoid calling scheduler.step() after the last training iteration is done

Reviewed By: wat3rBro

Differential Revision: D38605135

fbshipit-source-id: 87a55309bf6d1f7e598b567cc2372b00b8885c7c

d2go/runner/default_runner.py

@@ -576,7 +576,12 @@ class Detectron2GoRunner(BaseRunner):
             trainer.storage.put_scalar(
                 "lr", optimizer.param_groups[0]["lr"], smoothing_hint=False
             )
-            scheduler.step()
+            if trainer.iter < cfg.SOLVER.MAX_ITER - 1:
+                # Since scheduler.step() is called after the backward at each iteration,
+                # this will cause "where = 1.0" in the scheduler after the last interation,
+                # which will trigger "IndexError: list index out of range" in StepParamScheduler.
+                # See test_warmup_stepwithfixedgamma in vision/fair/detectron2/tests:test_scheduler for an example
+                scheduler.step()
             # Note: when precise BN is enabled, some checkpoints will have more precise
             # statistics than others, if they are saved immediately after eval.
             if comm.is_main_process():