switch to stable softmax when sampling

lucidrains · Dec 1, 2021 · 1ad3ab8 · 1ad3ab8
1 parent 2df5ff0
commit 1ad3ab8
Show file tree

Hide file tree

Showing 2 changed files with 3 additions and 2 deletions.
diff --git a/dalle_pytorch/dalle_pytorch.py b/dalle_pytorch/dalle_pytorch.py
@@ -10,6 +10,7 @@
 from dalle_pytorch import distributed_utils
 from dalle_pytorch.vae import OpenAIDiscreteVAE, VQGanVAE
 from dalle_pytorch.transformer import Transformer, DivideMax
+from dalle_pytorch.attention import stable_softmax
 
 # helpers
 
@@ -441,7 +442,7 @@ def generate_texts(
             logits = logits[:, -1, :]
 
             filtered_logits = top_k(logits, thres = filter_thres)
-            probs = F.softmax(filtered_logits / temperature, dim = -1)
+            probs = stable_softmax(filtered_logits / temperature, dim = -1)
             sample = torch.multinomial(probs, 1)
 
             text_tokens = torch.cat((text_tokens, sample), dim=-1)

diff --git a/setup.py b/setup.py
@@ -4,7 +4,7 @@
   name = 'dalle-pytorch',
   packages = find_packages(),
   include_package_data = True,
-  version = '1.1.4',
+  version = '1.1.5',
   license='MIT',
   description = 'DALL-E - Pytorch',
   author = 'Phil Wang',