bug found for guidance branch. RuntimeError: indices should be either on cpu or on the same device as the indexed tensor (cpu)

[BrandeisPatrick]

Error message:

File "/home/patrickli/working/DiGressGuidance/src/diffusion/noise_schedule.py", line 79, in get_alpha_bar
return self.alphas_bar[t_int.long()]
RuntimeError: indices should be either on cpu or on the same device as the indexed tensor (cpu)

Solution:

replace
return self.alphas_bar[t_int.long()]
with
return self.alphas_bar.to(t_int.device)[t_int.long()]