I want to process groups of rows in 2D-tensor and calculate simple statistics(eg. mean or max). I know how to do it step by step, but that’s an extremely slow approach for big tensors.
Is there a way to use fancy indexing or something similar to achieve this goal?
Example of what I am trying to do:
X = torch.randn(100, 10)

list_of_indexes = [
    [1, 4, 7, 10],
    [10, 60, 40, 10, 5, 3, 0],
    [6, 4]
]

for indexes in list_of_indexes:
    print(X[indexes].mean(dim=0))