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