fix device

lllyasviel · Mar 4, 2024 · d72347e · d72347e
1 parent 72139b0
commit d72347e
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/ldm_patched/modules/model_sampling.py b/ldm_patched/modules/model_sampling.py
@@ -74,8 +74,8 @@ def sigma_max(self):
         return self.sigmas[-1]
 
     def timestep(self, sigma):
-        log_sigma = sigma.log()
-        dists = log_sigma.to(self.log_sigmas.device) - self.log_sigmas[:, None]
+        log_sigma = sigma.log().to(self.log_sigmas.device)
+        dists = log_sigma - self.log_sigmas[:, None]
         low_idx = dists.ge(0).cumsum(dim=0).argmax(dim=0).clamp(max=self.log_sigmas.shape[0] - 2)
         high_idx = low_idx + 1
         low, high = self.log_sigmas[low_idx], self.log_sigmas[high_idx]