device mismatch (indices are on gpu)

[aiXander]

Getting this error:

│ /home/xander/Projects/cog/GitHub_repos/CelebBasis/ldm/models/diffusion/ddpm.py:1087 in p_losses  │
│                                                                                                  │
│   1084 │   │   loss_simple = self.get_loss(model_output, target, mean=False).mean([1, 2, 3])     │
│   1085 │   │   loss_dict.update({f'{prefix}/loss_simple': loss_simple.mean()})                   │
│   1086 │   │                                                                                     │
│ ❱ 1087 │   │   logvar_t = self.logvar[t].to(self.device)                                         │
│   1088 │   │   loss = loss_simple / torch.exp(logvar_t) + logvar_t                               │
│   1089 │   │   # loss = loss_simple / torch.exp(self.logvar) + self.logvar                       │
│   1090 │   │   if self.learn_logvar:                                                             │
╰──────────────────────────────────────────────────────────────────────────────────────────────────╯
RuntimeError: indices should be either on cpu or on the same device as the indexed tensor (cpu)

Changing logvar_t = self.logvar[t].to(self.device) to logvar_t = self.logvar[t.cpu()].to(self.device)

Solved this, but that feels like the wrong place to fix this bug...

[ygtxr1997]

I met the same bug on a new conda environment (with Pytorch 2.0). But I didn't meet this bug on the old conda envrionment (with Pytorch 1.12, 1.10). Maybe it is a bug reduced by a new Pytorch version, I'm not sure. Modifying the code as you said can fix this problem, which has been updated in the code.

[okaris]

I had the same issue today trying to upgrade torch for being able to work with newer gpus. I’ll share my solution in a few days.