From 100ea1158c5fe6dfafe844982c300fd12d8bda74 Mon Sep 17 00:00:00 2001 From: Katherine Crowson Date: Sat, 12 Aug 2023 18:19:12 +0000 Subject: [PATCH] Rename weights to c_weight --- k_diffusion/layers.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/k_diffusion/layers.py b/k_diffusion/layers.py index 46652ea..c7ce979 100644 --- a/k_diffusion/layers.py +++ b/k_diffusion/layers.py @@ -36,11 +36,11 @@ def get_scalings(self, sigma): def loss(self, input, noise, sigma, **kwargs): c_skip, c_out, c_in = [utils.append_dims(x, input.ndim) for x in self.get_scalings(sigma)] + c_weight = self.weighting(sigma) noised_input = input + noise * utils.append_dims(sigma, input.ndim) model_output = self.inner_model(noised_input * c_in, sigma, **kwargs) target = (input - c_skip * noised_input) / c_out - weights = self.weighting(sigma) - return (model_output - target).pow(2).flatten(1).mean(1) * weights + return (model_output - target).pow(2).flatten(1).mean(1) * c_weight def forward(self, input, sigma, **kwargs): c_skip, c_out, c_in = [utils.append_dims(x, input.ndim) for x in self.get_scalings(sigma)]