[FIXED] Calculating mean across all dimensions of a tensor except nth dimension

Issue

I have a torch tensor of dimension [7, 12, 12, 197, 197]. I would like to calculate mean such that resultant vector is of shape [7,1,1,197,1]. Is there any other way rather than using a for loop? Could you provide answer with benchmark on execution time?

Solution

Torch provides mean along with dim as argument. This dim can take tuple of integers.

import torch
tensor = torch.randn(4, 5, 6, 7)

dims_to_reduce = (1,3)
means = torch.mean(tensor, dim=dims_to_reduce, keepdim=True)
print(means.shape)
>>torch.Size([4, 1, 6, 1])

Answered By - MSS