(ACTUAL) Min-SNR Weighting Strategy: Fixed SNR calculation to authors implementation

This commit is contained in:
AI-Casanova
2023-03-23 12:34:49 +00:00
parent a3c7d711e4
commit 518a18aeff
5 changed files with 17 additions and 14 deletions

View File

@@ -381,7 +381,7 @@ def train(args):
loss = loss.mean([1, 2, 3])
if args.min_snr_gamma:
loss = apply_snr_weight(loss, latents, noisy_latents, args.min_snr_gamma)
loss = apply_snr_weight(loss, timesteps, noise_scheduler, args.min_snr_gamma)
loss_weights = batch["loss_weights"] # 各sampleごとのweight
loss = loss * loss_weights