Compute the weighted average in PyTorch

[sadra-barikbin]

Yes, that’s correct. To write it shorter:

weighted_average = (A@W)/W.sum()