Solved by KFrank in post #3 
              
              
                Hi Aesteban! 

The calculation you are asking for can be performed with a single 
matrix multiplication: 
>>> import torch
>>> torch.__version__
'1.9.0'
>>> _ = torch.manual_seed (2021)
>>> def avg_matrix_outer_products_v2(a):
...      x_dim = a.shape[0]
...      ourter_products = torch.outer(a[0], …