fix max_step of transformer-xl (PaddlePaddle#149)

wangxicoding · Mar 15, 2021 · 54b4e64 · 54b4e64
1 parent 18b0535
commit 54b4e64
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/examples/language_model/transformer-xl/train.py b/examples/language_model/transformer-xl/train.py
@@ -287,8 +287,8 @@ def do_train(args):
                     optimizer.set_lr(curr_lr)
             elif args.scheduler == 'noam':
                 scheduler.step()
-        if step_idx >= args.max_step:
-            break
+            if step_idx >= args.max_step:
+                return
 
     if args.save_model and rank == 0:
         model_dir = os.path.join(args.save_model, "step_final")