RuntimeError: indices should be either on cpu or on the same device as the indexed tensor (cpu)

[dorazhiyuyang]

---

RuntimeError Traceback (most recent call last)
Cell In[29], line 1
----> 1 x, n = sampling_fn(score_model)
2 show_samples(x)

File /workspace/pytorchcode/score_sde_pytorch-main/sampling.py:407, in get_pc_sampler..pc_sampler(model)
405 vec_t = torch.ones(shape[0], device=t.device) * t
406 x, x_mean = corrector_update_fn(x, vec_t, model=model)
--> 407 x, x_mean = predictor_update_fn(x, vec_t, model=model)
409 return inverse_scaler(x_mean if denoise else x), sde.N * (n_steps + 1)

File /workspace/pytorchcode/score_sde_pytorch-main/sampling.py:341, in shared_predictor_update_fn(x, t, sde, model, predictor, probability_flow, continuous)
339 else:
340 predictor_obj = predictor(sde, score_fn, probability_flow)
--> 341 return predictor_obj.update_fn(x, t)

File /workspace/pytorchcode/score_sde_pytorch-main/sampling.py:196, in ReverseDiffusionPredictor.update_fn(self, x, t)
195 def update_fn(self, x, t):
--> 196 f, G = self.rsde.discretize(x, t)
197 z = torch.randn_like(x)
198 x_mean = x - f

File /workspace/pytorchcode/score_sde_pytorch-main/sde_lib.py:104, in SDE.reverse..RSDE.discretize(self, x, t)
102 def discretize(self, x, t):
103 """Create discretized iteration rules for the reverse diffusion sampler."""
--> 104 f, G = discretize_fn(x, t)
105 rev_f = f - G[:, None, None, None] ** 2 * score_fn(x, t) * (0.5 if self.probability_flow else 1.)
106 rev_G = torch.zeros_like(G) if self.probability_flow else G

File /workspace/pytorchcode/score_sde_pytorch-main/sde_lib.py:251, in VESDE.discretize(self, x, t)
248 timestep = (t * (self.N - 1) / self.T).long()
249 sigma = self.discrete_sigmas.to(t.device)[timestep]
250 adjacent_sigma = torch.where(timestep == 0, torch.zeros_like(t),
--> 251 self.discrete_sigmas[timestep - 1].to(t.device))
252 f = torch.zeros_like(x)
253 G = torch.sqrt(sigma ** 2 - adjacent_sigma ** 2)

RuntimeError: indices should be either on cpu or on the same device as the indexed tensor (cpu)

[pace577]

Changing self.discrete_sigmas[timestep - 1].to(t.device) to self.discrete_sigmas.to(t.device)[timestep - 1] in this line of sde_lib.py seems to fix the problem.