Shuffling the indices as in https://github.com/noegroup/bgflow/blob/ae1dc6a9b3bcd7f0382c98614a572b2ae1b3a540/bgflow/distribution/sampling/dataset.py#L85 might be quite a bit slower than shuffling the tensors.