Efficient random voxel sampling with PyTorch

Question

I am trying to implement a fast sampling algorithm for point clouds. The idea is to overlay a regular grid over the point cloud and cluster all points within the same voxel. I do this with the voxel_grid function of Pytorch-Geometric. Then i want to randomly sample an arbitary number of points from each voxel. But i cannot come up with code that doesn't loop over each voxel and point cloud of the batch. Therefore, my algorithm is very slow.

What i have tried and worked well is to use scatter_min and scatter_max. This is very fast but i can only sample a single point each time. Additionaly, the selection is not random and invovles finding the min/max value of a specific feature in the given subset of points. So, how can i implement something like scatter_min/max but with random selection?