If you want to use a data loader (which efficiently extracts batches and uses multiprocessing), you can use the get_batch function on train_dataset defined in the code below.
Note that this assumes that an instance of train_data and train_labels are returned in the __getitem__(self, index) function of your dataset class.
from torch.utils import data 

class DynamicSampler(object):
    def __init__(self, max_size=100):
        self.next_batch = [0]
        self.max_size = max_size

    def select_sample(self, indList):
        self.next_batch = indList

    def __iter__(self):
        return iter(self.next_batch)

    def __len__(self):
        return self.max_size

def get_batch(dataset, indices=None, num_workers=2):
    sampler = DynamicSampler(len(indices))
    loader = data.DataLoader(dataset, 
                  batch_size=len(indices), 
                  sampler=sampler)
    
    sampler.select_sample(indices)

    return iter(loader).next()

if __name__ == '__main__':
    indices = np.arange(5, 100)
    batch = get_batch(train_dataset, indices)