Update generation_utils.py

generation_utils.py

@@ -425,7 +425,7 @@ class DreamGenerationMixin:
             t = timesteps[i]
             s = timesteps[i + 1]
         
-            logits[:,pad_token_id] += eos_penalty * torch.log(1-t+eps)
+            logits[:,:,pad_token_id] += eos_penalty * torch.log(1-t+eps)
 
             num_mask_token = mask_index.sum() / mask_index.shape[0]
             number_transfer_tokens = int(num_mask_token * (1 - s / t)) if i < steps - 1 else int(num_mask_token)
@@ -454,7 +454,7 @@ class DreamGenerationMixin:
                     hidden_states_, logits_ = self.forward_pass_drafter(x[active_indices], hidden_states[active_indices], num_transfered_)
                     hidden_states[active_indices] = hidden_states_
                     logits[active_indices] = logits_
-                    logits[:,pad_token_id] += eos_penalty * torch.log(1-t+eps)
+                    logits[:,:,pad_token_id] += eos_penalty * torch.log(1-t+eps)
 
             # this allows user-defined token control of the intermediate steps
             x = generation_tokens_hook_func(i, x, logits)