From 48301e389a793b9ccfdcb444c8b0acc43ff5a273 Mon Sep 17 00:00:00 2001
From: Gazzoo-byte <73721238+Gazzoo-byte@users.noreply.github.com>
Date: Fri, 23 Dec 2022 21:39:40 +0000
Subject: [PATCH] Fixes an issue when training embedding

In certain cases (most notably on Colab using Automatic1111's WebUI, but not limited to that example) trying to train an embedding would cause https://github.com/AUTOMATIC1111/stable-diffusion-webui/issues/5821#issue-1501645853 this error.
---
 ldm/models/diffusion/ddpm.py | 1 +
 1 file changed, 1 insertion(+)

diff --git a/ldm/models/diffusion/ddpm.py b/ldm/models/diffusion/ddpm.py
index 6090212..7bd6bb6 100644
--- a/ldm/models/diffusion/ddpm.py
+++ b/ldm/models/diffusion/ddpm.py
@@ -900,6 +900,7 @@ class LatentDiffusion(DDPM):
         loss_simple = self.get_loss(model_output, target, mean=False).mean([1, 2, 3])
         loss_dict.update({f'{prefix}/loss_simple': loss_simple.mean()})
 
+        self.logvar = self.logvar.to(self.device)
         logvar_t = self.logvar[t].to(self.device)
         loss = loss_simple / torch.exp(logvar_t) + logvar_t
         # loss = loss_simple / torch.exp(self.logvar) + self.logvar