Assuming we have a tensor A, with shape

```
CxH
```

And tensor W, representing the weights, with shape `H`

Tensor A represents some features with dim `C`

for each spatial location in `H`

(flattened), and weights `W`

represents the weight for each spatial location. Instead of computing the mean via:

```
torch.mul(A, W).mean(1)
```

How can we compute the weighted average ?

The output dim should be of size `C`

.

Would it be:

```
Z = torch.mul(A, W)
Weighted_average = torch.sum(Z, dim=1) / torch.sum(W)
```