From 97401adf961322a8276fa52d55b02b025c8cdba9 Mon Sep 17 00:00:00 2001
From: Phil Wang <lucidrains@gmail.com>
Date: Fri, 5 Feb 2021 19:31:07 -0800
Subject: [PATCH] fix final bug

---
 deep_daze/deep_daze.py | 7 ++++---
 deep_daze/version.py   | 2 +-
 2 files changed, 5 insertions(+), 4 deletions(-)

diff --git a/deep_daze/deep_daze.py b/deep_daze/deep_daze.py
index 9f1c324..e00f7db 100644
--- a/deep_daze/deep_daze.py
+++ b/deep_daze/deep_daze.py
@@ -131,7 +131,7 @@ def __init__(
 
         self.generate_size_schedule()
 
-    def forward(self, text, return_loss=True):
+    def forward(self, text, return_loss=True, dry_run=False):
         out = self.model()
         out = norm_siren_output(out)
 
@@ -153,7 +153,8 @@ def forward(self, text, return_loss=True):
             image_embed = perceptor.encode_image(image)
             text_embed = perceptor.encode_text(text)
 
-        self.num_batches_processed += self.batch_size
+        if not dry_run:
+            self.num_batches_processed += self.batch_size
 
         loss = -self.loss_coef * torch.cosine_similarity(text_embed, image_embed, dim=-1).mean()
         return loss
@@ -354,7 +355,7 @@ def forward(self):
 
         tqdm.write(f'Imagining "{self.text}" from the depths of my weights...')
 
-        self.model(self.encoded_text) # do one warmup step due to potential issue with CLIP and CUDA
+        self.model(self.encoded_text, dry_run = True) # do one warmup step due to potential issue with CLIP and CUDA
 
         if self.open_folder:
             open_folder('./')
diff --git a/deep_daze/version.py b/deep_daze/version.py
index 1e4826d..5bf52d5 100644
--- a/deep_daze/version.py
+++ b/deep_daze/version.py
@@ -1 +1 @@
-__version__ = '0.4.7'
+__version__ = '0.4.8'