From e967e786030bb4683d0879f906746d51966c44dd Mon Sep 17 00:00:00 2001
From: NouamaneTazi <nouamane98@gmail.com>
Date: Wed, 20 Nov 2024 08:59:26 +0000
Subject: [PATCH] .

---
 src/nanotron/sanity_checks.py | 4 ----
 1 file changed, 4 deletions(-)

diff --git a/src/nanotron/sanity_checks.py b/src/nanotron/sanity_checks.py
index 1c9c0ce3..650607da 100644
--- a/src/nanotron/sanity_checks.py
+++ b/src/nanotron/sanity_checks.py
@@ -232,10 +232,6 @@ def before_optim_step_sanity_checks(
                 msg=lambda err: f"[Before optimizer step] Tied weights {name} are not synchronized. {err}",
             )
 
-        # SANITY CHECK: optimizer's lr is synchronized with lr_scheduler
-        for group in unwrapped_model.optimizer.param_groups:
-            assert group["lr"] == unwrapped_model.lr_scheduler.get_last_lr()[0]
-
         # SANITY CHECK: run model specific sanity checks
         unwrapped_model.before_optim_step_sanity_checks()