From 2f636466cae51f995cb996682ca0dec6c563cfa2 Mon Sep 17 00:00:00 2001
From: wailovet <wailovet@163.com>
Date: Fri, 28 Jun 2024 01:09:54 +0800
Subject: [PATCH] 'update'

---
 hook_HYDiT_main.py | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/hook_HYDiT_main.py b/hook_HYDiT_main.py
index af85446..809bd4e 100644
--- a/hook_HYDiT_main.py
+++ b/hook_HYDiT_main.py
@@ -529,7 +529,7 @@ def Core(args, LOG):
 
     if args.use_fp16:
         optimizer = torch.optim.AdamW(
-            model.parameters(), lr=args.lr, weight_decay=args.weight_decay, eps=4e-6)
+            model.parameters(), lr=args.lr, weight_decay=args.weight_decay, eps=TRAIN_CONFIG.get("adam_epsilon", 1e-8))
     else:
         optimizer = torch.optim.AdamW(
             model.parameters(), lr=args.lr, weight_decay=args.weight_decay)
@@ -784,7 +784,7 @@ def easy_sample_images(
                                            requires_safety_checker=False,
                                            embedder_t5=embedder_t5,
                                            )
-        
+
         pipeline.to("cuda")
 
         style = torch.as_tensor([0, 0] * batch_size, device="cuda")