From f2e50a351d69a2825f454ba464dde0b7a400dbbe Mon Sep 17 00:00:00 2001 From: Akshay Kashyap Date: Wed, 29 May 2024 19:44:06 -0700 Subject: [PATCH] test condition --- tests/test_lr_schedulers.py | 5 ++--- 1 file changed, 2 insertions(+), 3 deletions(-) diff --git a/tests/test_lr_schedulers.py b/tests/test_lr_schedulers.py index f17cd7ce..0bfa503a 100644 --- a/tests/test_lr_schedulers.py +++ b/tests/test_lr_schedulers.py @@ -38,7 +38,7 @@ def test_cosine_lr_scheduler(): for i in range(num_training_steps): optimizer.step() lr_scheduler.step() - if i % intervals == 0: + if i == 0 or (i + 1) % intervals == 0: recorded = recorded_lrs_at_intervals.pop(0) assert math.isclose( lr_scheduler.get_last_lr()[0], recorded @@ -61,9 +61,8 @@ def test_inverse_sqrt_lr_scheduler(): ) for i in range(num_training_steps): - optimizer.step() lr_scheduler.step() - if i % intervals == 0: + if i == 0 or (i + 1) % intervals == 0: recorded = recorded_lrs_at_intervals.pop(0) assert math.isclose( lr_scheduler.get_last_lr()[0], recorded