diff options
Diffstat (limited to 'src/train.py')
| -rw-r--r-- | src/train.py | 3 |
1 files changed, 3 insertions, 0 deletions
diff --git a/src/train.py b/src/train.py index 2523325..8ce5394 100644 --- a/src/train.py +++ b/src/train.py @@ -105,6 +105,9 @@ def evaluate_gpu(model, gpu_ds: dict[str, torch.Tensor | int | bool], examples: o = model(x, y_in) nxt = o["logits"][:, -1].argmax(-1) y_in = torch.cat([y_in, nxt[:, None]], dim=1) + # If all sequences in the batch have emitted at least one EOS, we can exit safely + if ((y_in[:, 1:] == cfg.eos_id).any(dim=1)).all(): + break gen = y_in[:, 1:].cpu() for i, ex in enumerate(chunk_examples): |