pytorch 分割一個tensor並求平均

pytorch的torch.add()及torch.split()函數

import torch
# outputs是一個[batch, seq, 40]維的tensor，把outputs分割成兩個[batch, seq, 20]的tensor，並每個元素求平均值
add = torch.add(*torch.split(outputs, 20, dim=2)) / 2