Эта функция называется транслированием (broadcasting). batch_t формы (2, 3, 5, 5) умножается на unsqueezed_weights формы (3, 1, 1), в результате чего получается тензор формы (2, 3, 5, 5), в котором затем можно сложить третье измерение с конца (три канала):
# In[5]:
unsqueezed_weights = weights.unsqueeze(-1).unsqueeze_(-1)
img_weights = (img_t * unsqueezed_weights)
batch_weights = (batch_t * unsqueezed_weights)
img_gray_weighted = img_weights.sum(-3)
batch_gray_weighted = batch_weights.sum(-3)
batch_weights.shape, batch_t.shape, unsqueezed_weights.shape
# Out[5]:
(torch.Size([2, 3, 5, 5]), torch.Size([2, 3, 5, 5]), torch.Size([3, 1, 1]))
