burn_dataset/dataset/
iterator.rs

1use crate::dataset::{Dataset, DatasetError};
2use std::error::Error;
3use std::iter::Iterator;
4
5/// Dataset iterator.
6pub struct DatasetIterator<'a, I, E = DatasetError>
7where
8    E: Error + Send + Sync + 'static,
9{
10    current: usize,
11    dataset: &'a dyn Dataset<I, E>,
12    len: usize,
13}
14
15impl<'a, I, E> DatasetIterator<'a, I, E>
16where
17    E: Error + Send + Sync + 'static,
18{
19    /// Creates a new dataset iterator.
20    pub fn new<D>(dataset: &'a D) -> Self
21    where
22        D: Dataset<I, E>,
23    {
24        DatasetIterator {
25            current: 0,
26            dataset,
27            len: dataset.len(),
28        }
29    }
30}
31
32impl<I, E> Iterator for DatasetIterator<'_, I, E>
33where
34    E: Error + Send + Sync + 'static,
35{
36    type Item = Result<I, E>;
37
38    fn next(&mut self) -> Option<Result<I, E>> {
39        if self.current >= self.len {
40            return None;
41        }
42
43        let index = self.current;
44        self.current += 1;
45
46        Some(self.dataset.get(index))
47    }
48}
burn_dataset/dataset/iterator.rs

burn_dataset/dataset/
iterator.rs
