from torch.utils.data import IterableDataset

class CustomIterableDataset(IterableDataset):

    def __init__(self, filename):

        #Store the filename in object's memory
        self.filename = filename

        #And that's it, we no longer need to store the contents in the memory

    def __iter__(self):

        #Create an iterator
        file_itr = open(self.filename)

        return file_itr