diff options
Diffstat (limited to 'data/csv.py')
| -rw-r--r-- | data/csv.py | 1 |
1 files changed, 1 insertions, 0 deletions
diff --git a/data/csv.py b/data/csv.py index 233f5d8..619452e 100644 --- a/data/csv.py +++ b/data/csv.py | |||
| @@ -403,6 +403,7 @@ class VlpnDataset(IterableDataset): | |||
| 403 | if len(batch) >= batch_size: | 403 | if len(batch) >= batch_size: |
| 404 | yield batch | 404 | yield batch |
| 405 | batch = [] | 405 | batch = [] |
| 406 | continue | ||
| 406 | 407 | ||
| 407 | bucket_mask = mask.logical_and(self.bucket_assignments == bucket) | 408 | bucket_mask = mask.logical_and(self.bucket_assignments == bucket) |
| 408 | bucket_items = self.bucket_items[bucket_mask] | 409 | bucket_items = self.bucket_items[bucket_mask] |