title: Loading enormous custom dataset using IterableDataset
body: <p>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:
  <code>TypeError: iter() returned non-iterator of type 'TensorDataset'</code></p>
  <p>Following is the custom dataset class I wrote:</p>
  <pre><code>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
  </code></pre>
  <p>And I use it like this:</p>
  <pre><code>train_dataset=MyDataset(args)
  train_dataloader = DataLoader(train_dataset, batch_size=args.train_batch_size)
  </code></pre>
  <p>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?</p>