diff --git a/CODEOWNERS b/CODEOWNERS
index be96e068c15ae..9785fb3e55e58 100644
--- a/CODEOWNERS
+++ b/CODEOWNERS
@@ -5,7 +5,9 @@ orttraining/*.py @thiagocrepaldi @spandantiwari @BowenBao @liqunfu
 orttraining/orttraining/python/** @thiagocrepaldi @spandantiwari @BowenBao @liqunfu
 orttraining/orttraining/test/python/** @thiagocrepaldi @spandantiwari @BowenBao @liqunfu
 orttraining/pytorch_frontend_examples/** @thiagocrepaldi @spandantiwari @BowenBao @liqunfu
-onnxruntime/*.py @thiagocrepaldi @spandantiwari @BowenBao @liqunfu
-onnxruntime/python/** @thiagocrepaldi @spandantiwari @BowenBao @liqunfu
-onnxruntime/test/python/** @thiagocrepaldi @spandantiwari @BowenBao @liqunfu
+onnxruntime/python/training/** @thiagocrepaldi @spandantiwari @BowenBao @liqunfu
+onnxruntime/test/python/onnxruntime_test_ort_trainer.py @thiagocrepaldi @spandantiwari @BowenBao @liqunfu
+onnxruntime/test/python/onnxruntime_test_ort_trainer_with_mixed_precision.py @thiagocrepaldi @spandantiwari @BowenBao @liqunfu
+onnxruntime/test/python/onnxruntime_test_training_unit_tests.py @thiagocrepaldi @spandantiwari @BowenBao @liqunfu
+onnxruntime/test/python/onnxruntime_test_training_unittest_utils.py @thiagocrepaldi @spandantiwari @BowenBao @liqunfu
 samples/python/** @thiagocrepaldi @spandantiwari @BowenBao @liqunfu
diff --git a/orttraining/orttraining/python/experimental/optim/lr_scheduler.py b/orttraining/orttraining/python/experimental/optim/lr_scheduler.py
index 3583b363497c3..15d37efbc5cc3 100644
--- a/orttraining/orttraining/python/experimental/optim/lr_scheduler.py
+++ b/orttraining/orttraining/python/experimental/optim/lr_scheduler.py
@@ -232,7 +232,11 @@ class PolyWarmupLRScheduler(_LRScheduler):
     Learning rate update strategy:
         When current_step < warmup
             lr = base_lr * (current_step / max(1, num_warmup_steps))
+        When current_step > total_steps
+            lr = lr_end / lr
         Otherwise
+            lr =  decay / lr, where decay is
+                (lr - lr_end) * (1 - (current_step - num_warmup_steps) / (total_steps - num_warmup_steps)) ** power + lr_end
 
     Args:
         total_steps (int): total training steps for learning.