From 1f1cae6c5a1c55af9578ddfeda7dcfb0d1e1c2f3 Mon Sep 17 00:00:00 2001 From: ddPn08 Date: Thu, 1 Jun 2023 10:32:34 +0900 Subject: [PATCH] make the device of `snr_weight` the same as loss --- library/custom_train_functions.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/library/custom_train_functions.py b/library/custom_train_functions.py index f32f050e..fa24f9fa 100644 --- a/library/custom_train_functions.py +++ b/library/custom_train_functions.py @@ -14,7 +14,7 @@ def apply_snr_weight(loss, timesteps, noise_scheduler, gamma): all_snr = (alpha / sigma) ** 2 snr = torch.stack([all_snr[t] for t in timesteps]) 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 + snr_weight = torch.minimum(gamma_over_snr, torch.ones_like(gamma_over_snr)).float().to(loss.device) # from paper loss = loss * snr_weight return loss