Merge pull request #36 from lucasnewman/fix-cond-drop

Fix conditional drop for CFG when conditioning on semantic/phoneme tokens
lucidrains · Nov 22, 2023 · 79f6672 · 79f6672
2 parents ca8482b + c113619
commit 79f6672
Showing 1 changed file with 3 additions and 1 deletion.
diff --git a/voicebox_pytorch/voicebox_pytorch.py b/voicebox_pytorch/voicebox_pytorch.py
@@ -1023,6 +1023,8 @@ def forward(
 
         # classifier free guidance
 
+        cond_ids = cond_token_ids
+
         if cond_drop_prob > 0.:
             cond_drop_mask = prob_mask_like(cond.shape[:1], cond_drop_prob, self.device)
 
@@ -1043,7 +1045,7 @@ def forward(
         cond_emb = None
 
         if self.condition_on_text:
-            cond_emb = self.to_cond_emb(cond_token_ids)
+            cond_emb = self.to_cond_emb(cond_ids)
 
             cond_emb_length = cond_emb.shape[-2]
             if cond_emb_length != seq_len: