From 8bbc956ff14852e6faf9bc21bba884c34630626f Mon Sep 17 00:00:00 2001
From: Phil Wang <lucidrains@gmail.com>
Date: Wed, 31 Aug 2022 17:18:58 -0700
Subject: [PATCH] fix bug with misnamed variable in diffusion prior network

---
 dalle2_pytorch/dalle2_pytorch.py | 2 +-
 dalle2_pytorch/version.py        | 2 +-
 2 files changed, 2 insertions(+), 2 deletions(-)

diff --git a/dalle2_pytorch/dalle2_pytorch.py b/dalle2_pytorch/dalle2_pytorch.py
index ee543b1..98fe763 100644
--- a/dalle2_pytorch/dalle2_pytorch.py
+++ b/dalle2_pytorch/dalle2_pytorch.py
@@ -1070,7 +1070,7 @@ class DiffusionPriorNetwork(nn.Module):
 
         null_text_embeds = self.null_text_embeds.to(text_embed.dtype)
 
-        text_embeds = torch.where(
+        text_embed = torch.where(
             text_keep_mask,
             text_embed,
             null_text_embeds
diff --git a/dalle2_pytorch/version.py b/dalle2_pytorch/version.py
index 78a6e51..b436016 100644
--- a/dalle2_pytorch/version.py
+++ b/dalle2_pytorch/version.py
@@ -1 +1 @@
-__version__ = '1.10.3'
+__version__ = '1.10.4'