From cbbf077593bfb7770cfc0889b601acd915feca78 Mon Sep 17 00:00:00 2001 From: comfyanonymous Date: Wed, 18 Dec 2024 18:23:28 -0500 Subject: [PATCH] Small optimizations. --- comfy/ldm/modules/distributions/distributions.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/comfy/ldm/modules/distributions/distributions.py b/comfy/ldm/modules/distributions/distributions.py index f2b8ef90113..df987c5ec3f 100644 --- a/comfy/ldm/modules/distributions/distributions.py +++ b/comfy/ldm/modules/distributions/distributions.py @@ -30,10 +30,10 @@ def __init__(self, parameters, deterministic=False): self.std = torch.exp(0.5 * self.logvar) self.var = torch.exp(self.logvar) if self.deterministic: - self.var = self.std = torch.zeros_like(self.mean).to(device=self.parameters.device) + self.var = self.std = torch.zeros_like(self.mean, device=self.parameters.device) def sample(self): - x = self.mean + self.std * torch.randn(self.mean.shape).to(device=self.parameters.device) + x = self.mean + self.std * torch.randn(self.mean.shape, device=self.parameters.device) return x def kl(self, other=None):