Restore LoRA step scheduling functionality

2025-12-13 00:54:22 +02:00
parent 4a3ab6958a
commit 164a6bbebd
2 changed files with 7 additions and 4 deletions
@@ -264,6 +264,11 @@ class CustomLinear(nn.Linear):
        del weight, input, bias
        return out

+def update_lora_step(module, step):
+    for name, submodule in module.named_modules():
+        if isinstance(submodule, CustomLinear) and hasattr(submodule, "_step"):
+            submodule._step.fill_(step)
+
 def remove_lora_from_module(module):
    for name, submodule in module.named_modules():
        if hasattr(submodule, "lora_diffs"):