11 lines
259 B
Python
11 lines
259 B
Python
from torch.utils.data import IterableDataset
|
|
|
|
|
|
class TorchIterableDataset(IterableDataset):
|
|
def __init__(self, generator_func):
|
|
self.generator_func = generator_func
|
|
|
|
def __iter__(self):
|
|
it = self.generator_func()
|
|
yield from it
|