RuntimeError: indices should be either on cpu or on the same device as the indexed tensor (cpu)

Hello,

I am trying to train on custom dataset (where I have already prepared 1 - 1 image pairs and my seeds.json looks like this [["0000000", ["0"]], ["0000001", ["1"]], ... ) with 3x NVIDIA TITAN RTX 24GB. Initialization of all the models works fine but during validation sanity check I am getting this error:

...
[rank0]:   File "/code/instruct-pix2pix/./stable_diffusion/ldm/models/diffusion/ddpm_edit.py", line 892, in forward
[rank0]:     return self.p_losses(x, c, t, *args, **kwargs)
[rank0]:   File "/code/instruct-pix2pix/./stable_diffusion/ldm/models/diffusion/ddpm_edit.py", line 1043, in p_losses
[rank0]:     logvar_t = self.logvar[t].to(self.device)
[rank0]: RuntimeError: indices should be either on cpu or on the same device as the indexed tensor (cpu)

Do you know where it can come from? I did not change anything in source code. Just prepared the data and updated paths in train config.

Thanks in advance!