Error: "euler_a" and "lms" not working, tensors on different devices

Same error for "euler_a" and "lms" noise schedulers.

Traceback (most recent call last):
  File "/content/LECO/./train_lora.py", line 334, in <module>
    main(args)
  File "/content/LECO/./train_lora.py", line 321, in main
    train(config, prompts)
  File "/content/LECO/./train_lora.py", line 166, in train
    latents = train_util.get_initial_latents(
  File "/content/LECO/train_util.py", line 55, in get_initial_latents
    latents = noise * scheduler.init_noise_sigma
RuntimeError: Expected all tensors to be on the same device, but found at least two devices, cuda:0 and cpu!

"ddim" still works fine.

p1atdev / LECO

Error: "euler_a" and "lms" not working, tensors on different devices #20