fix final bug

lucidrains · Feb 6, 2021 · 97401ad · 97401ad
1 parent 91a9339
commit 97401ad
Show file tree

Hide file tree

Showing 2 changed files with 5 additions and 4 deletions.
diff --git 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
@@ -1 +1 @@
-__version__ = '0.4.7'
+__version__ = '0.4.8'