From 40128b7dc05a0bfbfee924e3c3b4efdbfba066d3 Mon Sep 17 00:00:00 2001
From: rockerBOO <rockerboo@gmail.com>
Date: Sat, 12 Apr 2025 04:10:48 -0400
Subject: [PATCH] Use args.wavelet_loss to activate

---
 train_network.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/train_network.py b/train_network.py
index 0f92e4f9..e2974c47 100644
--- a/train_network.py
+++ b/train_network.py
@@ -469,7 +469,7 @@ class NetworkTrainer:
         loss = train_util.conditional_loss(noise_pred.float(), target.float(), args.loss_type, "none", huber_c)
 
         wav_loss = None
-        if args.wavelet_loss_alpha:
+        if args.wavelet_loss:
             if args.wavelet_loss_rectified_flow:
                 # Calculate flow-based clean estimate using the target
                 flow_based_clean = noisy_latents - sigmas.view(-1, 1, 1, 1) * target