fix: Make sure timesteps_s in sample_tcd is on cpu

comfyanonymous · Apr 30, 2024 · 9b8f97c · 9b8f97c
1 parent 5cdf8cd
commit 9b8f97c
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/comfy/k_diffusion/sampling.py b/comfy/k_diffusion/sampling.py
@@ -825,7 +825,7 @@ def sample_tcd(
     s_in = x.new_ones([x.shape[0]])
 
     model_sampling = model.inner_model.model_patcher.get_model_object("model_sampling")
-    timesteps_s = torch.floor((1 - eta) * model_sampling.timestep(sigmas)).to(dtype=torch.long).detach()
+    timesteps_s = torch.floor((1 - eta) * model_sampling.timestep(sigmas)).to(dtype=torch.long).detach().cpu()
     timesteps_s[-1] = 0
     alpha_prod_s = model_sampling.alphas_cumprod[timesteps_s]
     beta_prod_s = 1 - alpha_prod_s