def METHOD_NAME(self) -> DataLoader:
    train_dataset = data.WikiTextDataset(
        self.corpus,
        batch_size=self.context.get_per_slot_batch_size(),
    )
    batch_samp = data.BatchSamp(train_dataset, self.bptt)
    return DataLoader(train_dataset, batch_sampler=batch_samp)