huggingface/candle
1
0
Fork 0
mirror of https://github.com/huggingface/candle.git synced 2025-06-15 18:28:24 +00:00
Code Issues Packages Projects Releases Wiki Activity

Fix Batcher iterator break when return_last_incomplete_batch and items.is_empty (#2654) (#2655)

Browse Source
This commit is contained in:
hhllhhyyds
2024-12-24 15:41:26 +08:00
committed by GitHub
parent 1be6b090c7
commit 11aa30be10
1 changed files with 4 additions and 4 deletions
Download Patch File Download Diff File Expand all files Collapse all files

8
candle-datasets/src/batcher.rs
View File

@ -78,7 +78,7 @@ impl<I: Iterator<Item = Tensor>> Iterator for Batcher<Iter1<I>> {
match self.inner.inner.next() {
Some(item) => items.push(item),
None => {
if self.return_last_incomplete_batch {
if self.return_last_incomplete_batch && !items.is_empty() {
break;
}
return None;
@ -102,7 +102,7 @@ impl<I: Iterator<Item = (Tensor, Tensor)>> Iterator for Batcher<Iter2<I>> {
ys.push(y)
}
None => {
if self.return_last_incomplete_batch {
if self.return_last_incomplete_batch && !xs.is_empty() && !ys.is_empty() {
break;
}
return None;
@ -127,7 +127,7 @@ impl<I: Iterator<Item = Result<Tensor>>> Iterator for Batcher<IterResult1<I>> {
match self.inner.inner.next() {
Some(item) => items.push(item),
None => {
if self.return_last_incomplete_batch {
if self.return_last_incomplete_batch && !items.is_empty() {
break;
}
return None;
@ -154,7 +154,7 @@ impl<I: Iterator<Item = Result<(Tensor, Tensor)>>> Iterator for Batcher<IterResu
}
Some(Err(err)) => errs.push(err),
None => {
if self.return_last_incomplete_batch {
if self.return_last_incomplete_batch && !xs.is_empty() && !ys.is_empty() {
break;
}
return None;