Fix the early stopping callback #559
This commit is contained in:
parent
a1f12d607f
commit
8c8e8b4450
1 changed files with 1 additions and 1 deletions
|
@ -25,7 +25,7 @@ class _SentinelTokenStoppingCriteria(transformers.StoppingCriteria):
|
||||||
if trimmed_sample.shape[-1] < self.sentinel_token_ids[i].shape[-1]:
|
if trimmed_sample.shape[-1] < self.sentinel_token_ids[i].shape[-1]:
|
||||||
continue
|
continue
|
||||||
for window in trimmed_sample.unfold(0, self.sentinel_token_ids[i].shape[-1], 1):
|
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 True
|
||||||
return False
|
return False
|
||||||
|
|
||||||
|
|
Loading…
Add table
Add a link
Reference in a new issue