Loading enormous custom dataset using IterableDataset
I have a huge dataset with features (input_id,input_mask,segment_id,label_id) saved in batches of 64 in a pickle file. I read this file, create a TensorDataset and pass to dataloader for training. Since the features file is too big to create the full TensorDataset, I want to convert TensorDataset to IterableDataset so that one batch of samples can be retrieved from the features file at a time and passed on to the dataloader. But while training, I get the following error:
TypeError: iter() returned non-iterator of type 'TensorDataset'
Following is the custom dataset class I wrote:
class MyDataset(IterableDataset):

    def __init__(self,args):
        self.args=args
       
    def get_features(self,filename):
        with open(filename, &quot;rb&quot;) as f:
            while True:
                try:
                    yield pickle.load(f)
                except EOFError:
                    break  
                    
    def process(self,args):
        if args.cached_features_file:
            cached_features_file = args.cached_features_file

        if os.path.exists(cached_features_file):
            features=self.get_features(cached_features_file)

        feat = next (features)
        li=list(feat)
        all_input_ids=torch.tensor([f.input_ids for f in li ], dtype=torch.long)
        all_input_mask= torch.tensor([f.input_mask for f in li ], dtype=torch.long)
        all_segment_ids= torch.tensor([f.segment_ids for f in li], dtype=torch.long)
        all_label_ids = torch.tensor([f.label_id for f in li ], dtype=torch.long)
        
        dataset = TensorDataset(all_input_ids, all_input_mask, all_segment_ids, all_label_ids)
        return dataset
      
    def __iter__(self):
        dataset=self.process(self.args)       
        return dataset

And I use it like this:
train_dataset=MyDataset(args)
train_dataloader = DataLoader(train_dataset, batch_size=args.train_batch_size)

I understand that TensorDataset is map-style requiring index while IterableDataset is iterable-style which is the reason for error. Even if I return a list/tuple of feature tensors instead of TensorDataset, I get similar error. Can someone please tell me how to load a batched dataset in the correct way with IterableDataset?

I solved the problem by saving the dataset in a different manner. I saved the features as dictionary objects incrementally pickled in a pickle file and just read them off one at a time and pass on to dataloader for processing. Batching is done automatically by the dataloader. This is how the custom class looks now:
class MyDataset(IterableDataset):

    def __init__(self,filename):
     
        self.filename=filename
        super().__init__()
                    
    def process(self,filename):
        with open(filename, &quot;rb&quot;) as f:
            while True:
                try:
                    yield pickle.load(f)
                except EOFError:
                    break

    def __iter__(self):
        dataset=self.process(self.filename)          
        return dataset

