fix bug in MSE color mean loss

ZERONE182 · May 19, 2023 · 4257c98 · 4257c98
1 parent 2ea23c9
commit 4257c98
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/utils.py b/utils.py
@@ -69,7 +69,7 @@ def p_losses(denoise_model, img, R, T, K, logsnr, noise=None, loss_type="l2", co
     rec_color_mean = torch.mean(rec_img, dim=(2, 3))
 
     color_loss = torch.nn.MSELoss()
-    color_loss = (logsnr.to(dev()) + 20) / 20 * color_loss(img_color_mean.to(dev()), rec_color_mean)
+    color_loss = (logsnr.to(dev()) + 20) / 20 * color_loss(img_color_mean.to(dev()), rec_color_mean).mean()
 
     return loss + color_loss