A bug in device during sampling

[Z7Gao]

Hi, I encountered a bug when running the real image sampling script.

File "xxxxx/score-MRI/sde_lib.py", line 156, in discretize self.discrete_sigmas[timestep - 1].to(t.device)) RuntimeError: indices should be either on cpu or on the same device as the indexed tensor (cpu)

a simple fix can solve this problem

adjacent_sigma = torch.where(timestep == 0, torch.zeros_like(t),
                                 self.discrete_sigmas[timestep.item() - 1].to(t.device))