| 73 | |
| 74 | |
| 75 | class TrainDataLoaderIter(DataLoaderIter): |
| 76 | |
| 77 | def __init__( |
| 78 | self, data_loader: DataLoader, image_extractor: Callable, label_extractor: Callable, auto_reset: bool = True |
| 79 | ) -> None: |
| 80 | super().__init__(data_loader, image_extractor, label_extractor) |
| 81 | self.auto_reset = auto_reset |
| 82 | |
| 83 | def __next__(self): |
| 84 | try: |
| 85 | batch = next(self._iterator) |
| 86 | inputs, labels = self.inputs_labels_from_batch(batch) |
| 87 | except StopIteration: |
| 88 | if not self.auto_reset: |
| 89 | raise |
| 90 | self._iterator = iter(self.data_loader) |
| 91 | batch = next(self._iterator) |
| 92 | inputs, labels = self.inputs_labels_from_batch(batch) |
| 93 | |
| 94 | return inputs, labels |
| 95 | |
| 96 | |
| 97 | class ValDataLoaderIter(DataLoaderIter): |
no outgoing calls
no test coverage detected
searching dependent graphs…