fix bug in MSE color mean loss

ZERONE182 · May 19, 2023 · 9c4f8a7 · 9c4f8a7
1 parent c94187d
commit 9c4f8a7
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/utils.py b/utils.py
@@ -64,7 +64,7 @@ def p_losses(denoise_model, img, R, T, K, logsnr, noise=None, loss_type="l2", co
     else:
         raise NotImplementedError()
 
-    rec_img = reconstruct_z_start(z_noisy, predicted_noise, logsnr.to(dev()))
+    rec_img = reconstruct_z_start(z_noisy.to(dev()), predicted_noise, logsnr.to(dev()))
     img_color_mean = torch.mean(z, dim=(2, 3))
     rec_color_mean = torch.mean(rec_img, dim=(2, 3)).to(dev())