tweak

ku-nlp · Jul 18, 2024 · a338e9c · a338e9c
1 parent c78cb10
commit a338e9c
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/src/kwja/modules/components/logits_processor.py b/src/kwja/modules/components/logits_processor.py
@@ -103,7 +103,7 @@ def get_mask(
             for i, (prev_input_ids, logits) in enumerate(zip(batch_prev_input_ids.tolist(), batch_logits)):
                 batch_idx = i // self.num_beams
                 if self.is_finished[batch_idx]:
-                    mask = torch.ones_like(logits, dtype=torch.bool)
+                    mask = torch.zeros_like(logits, dtype=torch.bool)
                     mask[self.tokenizer.eos_token_id] = True
                     batch_masks.append(mask)
                     continue