Hello, I have a tensor O_k of shape [M,4,2]. I need to get an average across M for these tensors element-wise to obtain ave_O_k of shape [4,2]. I initially tried sum(O_k)/len(O_k), but this is very inefficient, I imagine it would be better to use torch’s torch.mean(O_k), but this function seems to take the row-wise average of O_k.
Have you tried doing torch.mean(O_k, dim=0)?
x = torch.randn(1000,4,2)
mean = torch.mean(x, dim=0)
mean.shape() #returns torch.Size([4, 2])