DreamArtist Web UI Extension Error: indices should be either on cpu or on the same device as the indexed tensor (cpu)

DreamArtist is an extension for Automatic1111’s stable diffusion web UI. It allows you to train embeddings by using one photo. I was excited to learn about this extension and wanted to give it a try.

However, when I tried it the first time, I got this error:

Traceback (most recent call last):
  File "/mnt/data/stable-diffusion-webui/modules/call_queue.py", line 45, in f
    res = list(func(*args, **kwargs))
  File "/mnt/data/stable-diffusion-webui/modules/call_queue.py", line 28, in f
    res = func(*args, **kwargs)
  File "/mnt/data/stable-diffusion-webui/extensions/DreamArtist-sd-webui-extension/scripts/dream_artist/ui.py", line 32, in train_embedding
    embedding, filename = dream_artist.cptuning.train_embedding(*args)
  File "/mnt/data/stable-diffusion-webui/extensions/DreamArtist-sd-webui-extension/scripts/dream_artist/cptuning.py", line 535, in train_embedding
    output = shared.sd_model(x, c_in, scale=(cfg_l, cfg_h), att_mask=None, dy_cfg_f=dy_cfg_f)
  File "/mnt/data/stable-diffusion-webui/venv/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1194, in _call_impl
    return forward_call(*input, **kwargs)
  File "/mnt/data/stable-diffusion-webui/repositories/stable-diffusion-stability-ai/ldm/models/diffusion/ddpm.py", line 846, in forward
    return self.p_losses(x, c, t, *args, **kwargs)
  File "/mnt/data/stable-diffusion-webui/extensions/DreamArtist-sd-webui-extension/scripts/dream_artist/cptuning.py", line 309, in p_losses_hook
    logvar_t = self.logvar[t_raw].to(self.device)
RuntimeError: indices should be either on cpu or on the same device as the indexed tensor (cpu)

At first glance, this looks daunting. Fortunately, Google search pointed me to the right direction to fix this problem. This YouTube video shows a similar problem. For my case, the file is different. I have to edit this file under stable diffusion web UI directory:

./extensions/DreamArtist-sd-webui-extension/scripts/dream_artist/cptuning.py

I need to find this line:

logvar_t = self.logvar[t_raw].to(self.device)

And insert this line before the logvar_t line:

self.logvar = self.logvar.to(self.device)

Note that the spaces are important, so make sure both lines have the same number of spaces in the front.

This is how it looks like after editing.

Save the file and you should be able to use the extension.