Fix more tests.

facebookresearch · Aug 30, 2020 · 125a5b4 · 125a5b4
1 parent abc3f2d
commit 125a5b4
Show file tree

Hide file tree

Showing 3 changed files with 4 additions and 3 deletions.
diff --git a/parlai/nn/lr_scheduler.py b/parlai/nn/lr_scheduler.py
@@ -60,7 +60,7 @@ def _init_warmup_scheduler(self, optimizer, states):
             self.warmup_scheduler = optim.lr_scheduler.LambdaLR(
                 optimizer, self._warmup_lr
             )
-            if 'warmup_scheduler' in states:
+            if states.get('warmup_scheduler'):
                 self.warmup_scheduler.load_state_dict(states['warmup_scheduler'])
         else:
             self.warmup_scheduler = None
@@ -307,7 +307,7 @@ def step(self, num_steps):
 
         Override this method to override the behavior for training schedulers.
         """
-        self._number_training_updates += 1
+        self._number_training_updates = num_steps
         if self._is_lr_warming_up():
             self.warmup_scheduler.step()
         else:

diff --git a/tests/test_lr_schedulers.py b/tests/test_lr_schedulers.py
@@ -96,7 +96,7 @@ def test_invsqrt(self):
         steps = self._run_pass(
             lr_scheduler='invsqrt', warmup_updates=50, invsqrt_lr_decay_gamma=1
         )
-        self.assertAlmostEquals(steps[-1], 0.0324272)
+        self.assertAlmostEquals(steps[-1], 0.0324443)
 
         # decay very slowly
         steps = self._run_pass(

diff --git a/tests/test_transformers.py b/tests/test_transformers.py
@@ -797,6 +797,7 @@ def _test_learning_rate_resuming(self, args):
                     init_model=os.path.join(tmpdir, 'model'),
                     model_file=os.path.join(tmpdir, 'newmodel2'),
                     lr_scheduler='reduceonplateau',
+                    log_every_n_secs=0.001,
                     **args,
                 )
             )