Option to change optimizer eps

NKI-AI · Dec 12, 2023 · 0d49cb8 · 0d49cb8
1 parent 1f829dd
commit 0d49cb8
Show file tree

Hide file tree

Showing 2 changed files with 2 additions and 0 deletions.
diff --git a/direct/config/defaults.py b/direct/config/defaults.py
@@ -52,6 +52,7 @@ class TrainingConfig(BaseConfig):
     optimizer: str = "Adam"
     lr: float = 5e-4
     weight_decay: float = 1e-6
+    eps: float = 1e-8
     batch_size: int = 2
 
     # LR Scheduler

diff --git a/direct/train.py b/direct/train.py
@@ -237,6 +237,7 @@ def setup_train(
         optimizer_params,
         lr=env.cfg.training.lr,
         weight_decay=env.cfg.training.weight_decay,
+        eps=env.cfg.training.eps,
     )  # noqa
 
     # Build the LR scheduler, we use a fixed LR schedule step size, no adaptive training schedule.