Or maybe this one
class MySampler(torch.utils.data.Sampler):
    def __init__(self, end_idx, seq_length):
        indices = []
        for i in range(len(end_idx) - 1):
            start = end_idx[i]
            end = end_idx[i + 1] - seq_length
            if start > end:
                indices.append(torch.arange(end, start))
            else:
                indices.append(torch.arange(start, end))
        indices = torch.cat(indices)
        self.indices = indices