Skip to main content

entrenar/prune/data_loader/
iter.rs

1//! Iterator over calibration batches.
2
3use super::loader::CalibrationDataLoader;
4use crate::train::Batch;
5
6/// Iterator over calibration batches.
7pub struct CalibrationDataIter<'a> {
8    loader: &'a CalibrationDataLoader,
9    position: usize,
10}
11
12impl<'a> CalibrationDataIter<'a> {
13    /// Create a new iterator over calibration batches.
14    pub(crate) fn new(loader: &'a CalibrationDataLoader) -> Self {
15        Self { loader, position: 0 }
16    }
17}
18
19impl<'a> Iterator for CalibrationDataIter<'a> {
20    type Item = &'a Batch;
21
22    fn next(&mut self) -> Option<Self::Item> {
23        let batch = self.loader.get_batch(self.position)?;
24        self.position += 1;
25        Some(batch)
26    }
27
28    fn size_hint(&self) -> (usize, Option<usize>) {
29        let remaining = self.loader.num_batches().saturating_sub(self.position);
30        (remaining, Some(remaining))
31    }
32}
33
34impl ExactSizeIterator for CalibrationDataIter<'_> {}