Small optimizations.

comfyanonymous · Dec 18, 2024 · cbbf077 · cbbf077
1 parent 0c04a6a
commit cbbf077
Showing 1 changed file with 2 additions and 2 deletions.
diff --git 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):