Three dimensional input for DataLoader


I have the following set-up:
I have a 3D data tensor with the axes (different_samples, timeseries_length, feature_dimensionality), i.e.,
the first dimension is the number of different timeseries samples I have, the second dimension is the length of the timeseries’ and the third dimension is the dimensionality of each timestep.
I want to plug this 3D tensor into the Dataloader. I have written my own BatchSampler, however, only for a 2D tensor containing a single timeseries.

What I want to do now, however, is that my custom BatchSampler first selects an index to choose one of the timeseries in the 3D tensor (i.e., an index for the first dimension) and then selects a starting index for the batch (i.e., an index for the second dimension) within the timeseries that was selected at first, which means that the BatchSampler has to return a 2D object.

My BatchSampler (which is still only able to handel 2D data input) currently returns a list of indices (I think the Sampler classes always have to return a list). And as far as I understand, this list of indices is then passed to the getitem method of the dataset object. So for my new set-up with a 3D data tensor, I would now require my BatchSampler to return something like a 2D object: The first dimension providing the index that selects the specific timeseries, and the second index dimension providing the indices that select the batch from that chosen timeseries. However, there are no 2D lists…
So how would I do that?
Or is there a completely different way to do so?
Help is very much appreciated!