I have an iterable dataset (and dataloader for it) which works fine. But after splitting it using torch.utils.data.random_split, and making the split datasets into a dataloader I’m getting a NotImplementedError.
Here’s a small example to reproduce the error:
class TestDataset(torch.utils.data.IterableDataset):
    def __init__(self):
        super().__init__()
        self.n = 10
        
    def __len__(self):
        return self.n
    
    def __iter__(self):
        for i in range(self.n):
            yield i
            
test_dataset = TestDataset()
test_dataset1, test_dataset2 = torch.utils.data.random_split(test_dataset, lengths=[8, 2])
test_dataloader1 = torch.utils.data.DataLoader(test_dataset1, batch_size=2)
next(iter(test_dataloader1))  # NotImplementedError

Why am I getting this error and how can I split an Iterable style dataset then?