Can't train at all in Colab, said there are tensors on two devices. Error when "Train with reconstruction" is ticked.

[xITmasterx]

I can't seem to run DreamArtist in a fresh install because of this error, and I can't seem to find out where did it went wrong. Got any ideas on how to fix this? Thanks in advance.

Traceback (most recent call last):
  File "/content/stable-diffusion-webui/modules/call_queue.py", line 45, in f
    res = list(func(*args, **kwargs))
  File "/content/stable-diffusion-webui/modules/call_queue.py", line 28, in f
    res = func(*args, **kwargs)
  File "/content/stable-diffusion-webui/extensions/DreamArtist-sd-webui-extension/scripts/dream_artist/ui.py", line 30, in train_embedding
    embedding, filename = dream_artist.cptuning.train_embedding(*args)
  File "/content/stable-diffusion-webui/extensions/DreamArtist-sd-webui-extension/scripts/dream_artist/cptuning.py", line 542, in train_embedding
    loss.backward()
  File "/usr/local/envs/automatic/lib/python3.10/site-packages/torch/_tensor.py", line 396, in backward
    torch.autograd.backward(self, gradient, retain_graph, create_graph, inputs=inputs)
  File "/usr/local/envs/automatic/lib/python3.10/site-packages/torch/autograd/__init__.py", line 173, in backward
    Variable._execution_engine.run_backward(  # Calls into the C++ engine to run the backward pass
RuntimeError: Expected all tensors to be on the same device, but found at least two devices, cuda:0 and cpu! (when checking argument for argument weight in method wrapper__convolution_backward)

[xITmasterx]

More details with the problem: Somehow this error is thrown out whenever I enabled the "Train with reconstruction" option.