From 8c72f4df00d47b6e5c7da1ed752993a318569de0 Mon Sep 17 00:00:00 2001
From: Rasaboun <40967731+Rasaboun@users.noreply.github.com>
Date: Fri, 7 Aug 2026 10:17:19 +0200
Subject: [PATCH] fix: make constant_with_warmup tolerate a missing total_iters

Same defect as the step branch, opposite direction: the branch used an
unguarded del kwargs['total_iters'], so get_lr_scheduler raised
KeyError: 'total_iters' whenever it was called without the trainer's
injection. Switch to the same pop() used elsewhere so every branch is
safe both with and without the key.
---
 toolkit/scheduler.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/toolkit/scheduler.py b/toolkit/scheduler.py
index 841f1244..2404ae69 100644
--- a/toolkit/scheduler.py
+++ b/toolkit/scheduler.py
@@ -43,7 +43,7 @@ def get_lr_scheduler(
         if 'num_warmup_steps' not in kwargs:
             print(f"WARNING: num_warmup_steps not in kwargs. Using default value of 1000")
             kwargs['num_warmup_steps'] = 1000
-        del kwargs['total_iters']
+        kwargs.pop('total_iters', None)
         return get_constant_schedule_with_warmup(optimizer, **kwargs)
     else:
         # try to use a diffusers scheduler