def METHOD_NAME(self) -> torch.Tensor:
    data = self.corpus.val if self.valid else self.corpus.train
    # Work out how cleanly we can divide the dataset into bsz parts.
    nbatch = data.size(0) // self.batch_size
    # Trim off any extra elements that wouldn't cleanly fit (remainders).
    data = data.narrow(0, 0, nbatch * self.batch_size)
    # Evenly divide the data across the bsz batches.
    data = data.view(self.batch_size, -1).t().contiguous()
    return data