In Python 3.6 ThreadPoolExecuter.map() is no more lazy. So in
with ThreadPoolExecutor(max_workers=self.num_workers) as e:
for batch in e.map(self.get_batch, iter(self.batch_sampler)):
yield get_tensor(batch, self.pin_memory)
e.map(self.get_batch, iter(self.batch_sampler)) will eagerly apply self.get_batch() to every element in the iterable iter(self.batch_sampler), store the result internally in an iterable and then return a generator wrapped around the iterable. In our case it means all batches will be processed and kept in memory simultaneously, even though we only want one batch to be generated at a time.
