diff --git a/modelscope/models/multi_modal/diffusion/diffusion.py b/modelscope/models/multi_modal/diffusion/diffusion.py
index de944498..04203e18 100644
--- a/modelscope/models/multi_modal/diffusion/diffusion.py
+++ b/modelscope/models/multi_modal/diffusion/diffusion.py
@@ -40,6 +40,7 @@ def discretized_gaussian_log_likelihood(x0, mean, log_scale):
 
 
 def _i(tensor, t, x):
+    tensor = tensor.to(x.device)
     shape = (x.size(0), ) + (1, ) * (x.ndim - 1)
     return tensor[t].view(shape).to(x)