From 0d49cb8bd9eefefbd4fbcaca5cd78e12592ff7bb Mon Sep 17 00:00:00 2001
From: georgeyiasemis <georgeyiasemis@hotmail.com>
Date: Wed, 13 Dec 2023 00:31:06 +0100
Subject: [PATCH] Option to change optimizer eps

---
 direct/config/defaults.py | 1 +
 direct/train.py           | 1 +
 2 files changed, 2 insertions(+)

diff --git a/direct/config/defaults.py b/direct/config/defaults.py
index 862f317d..52fd5433 100644
--- 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
index 33f4a73d..c994eb74 100644
--- 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.