From 4a1f3aba316c739cf2e917334018cecd42a79bba Mon Sep 17 00:00:00 2001 From: James Betker Date: Tue, 21 Jun 2022 20:12:54 -0600 Subject: [PATCH] come on guys... :(( --- codes/models/diffusion/gaussian_diffusion.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/codes/models/diffusion/gaussian_diffusion.py b/codes/models/diffusion/gaussian_diffusion.py index 472b210b..9c17e951 100644 --- a/codes/models/diffusion/gaussian_diffusion.py +++ b/codes/models/diffusion/gaussian_diffusion.py @@ -870,7 +870,7 @@ class GaussianDiffusion: s_err = (target - model_output) ** 2 if channel_balancing_fn is not None: s_err = channel_balancing_fn(s_err) - terms["mse_by_batch"] = s_err.view(s_err.shape[0], -1).mean(dim=1) + terms["mse_by_batch"] = s_err.reshape(s_err.shape[0], -1).mean(dim=1) terms["mse"] = mean_flat(s_err) terms["x_start_predicted"] = x_start_pred if "vb" in terms: