Merge pull request #1830 from nhamanasu/sd3

Added RAdamScheduleFree support
kohya-ss · Dec 15, 2024 · f2d38e6 · f2d38e6
2 parents d3305f9 + 8e378cf
commit f2d38e6
Showing 1 changed file with 5 additions and 1 deletion.
diff --git a/library/train_util.py b/library/train_util.py
@@ -4887,7 +4887,11 @@ def get_optimizer(args, trainable_params) -> tuple[str, str, object]:
             import schedulefree as sf
         except ImportError:
             raise ImportError("No schedulefree / schedulefreeがインストールされていないようです")
-        if optimizer_type == "AdamWScheduleFree".lower():
+
+        if optimizer_type == "RAdamScheduleFree".lower():
+            optimizer_class = sf.RAdamScheduleFree
+            logger.info(f"use RAdamScheduleFree optimizer | {optimizer_kwargs}")
+        elif optimizer_type == "AdamWScheduleFree".lower():
             optimizer_class = sf.AdamWScheduleFree
             logger.info(f"use AdamWScheduleFree optimizer | {optimizer_kwargs}")
         elif optimizer_type == "SGDScheduleFree".lower():