From 92ca1d5962d8e3468057ff132d858fed6acac9dc Mon Sep 17 00:00:00 2001
From: AI-Casanova <54461896+AI-Casanova@users.noreply.github.com>
Date: Tue, 21 Mar 2023 20:38:27 -0500
Subject: [PATCH] Min-SNR Weighting Strategy: Fixed SNR calculation to authors
 implementation

---
 library/custom_train_functions.py | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/library/custom_train_functions.py b/library/custom_train_functions.py
index f60ec7436..5e880c9a7 100644
--- a/library/custom_train_functions.py
+++ b/library/custom_train_functions.py
@@ -10,8 +10,8 @@ def apply_snr_weight(loss, latents, noisy_latents, gamma):
   gamma_over_snr = torch.div(torch.ones_like(snr)*gamma,snr)
   snr_weight = torch.minimum(gamma_over_snr,torch.ones_like(gamma_over_snr)).float() #from paper
   loss = loss * snr_weight
-  print(snr_weight)
+  #print(snr_weight)
   return loss
 
 def add_custom_train_arguments(parser: argparse.ArgumentParser):
-  parser.add_argument("--min_snr_gamma", type=float, default=0, help="gamma for reducing the weight of high loss timesteps. Lower numbers have stronger effect. 5 is recommended by paper.")
\ No newline at end of file
+  parser.add_argument("--min_snr_gamma", type=float, default=0, help="gamma for reducing the weight of high loss timesteps. Lower numbers have stronger effect. 5 is recommended by paper.")