This warmup-then-cosine schedule was designed for a 1000-STEP run — but the counter t advances once per EPOCH. After all 10 epochs (1000 real training steps) the lr has crawled to 0.001: the run spends its entire life in the first 1% of warmup and the cosine never happens. Fix the counter to advance once per step.
The break is on line 15 — but read the whole snippet first.