From 8c8e8b44508972a37fd15d760f9e4214e5105306 Mon Sep 17 00:00:00 2001 From: oobabooga <112222186+oobabooga@users.noreply.github.com> Date: Sat, 25 Mar 2023 12:35:52 -0300 Subject: [PATCH] Fix the early stopping callback #559 --- modules/callbacks.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/modules/callbacks.py b/modules/callbacks.py index 2ae9d908..8d30d615 100644 --- a/modules/callbacks.py +++ b/modules/callbacks.py @@ -25,7 +25,7 @@ class _SentinelTokenStoppingCriteria(transformers.StoppingCriteria): if trimmed_sample.shape[-1] < self.sentinel_token_ids[i].shape[-1]: continue for window in trimmed_sample.unfold(0, self.sentinel_token_ids[i].shape[-1], 1): - if torch.all(torch.eq(self.sentinel_token_ids[i], window)): + if torch.all(torch.eq(self.sentinel_token_ids[i][0], window)): return True return False