如何在 PyTorch 中处理批量数据？

Question

我有一个函数可以根据数据的几个特征计算一个新特征：

def my_func(f1, f2):
``
f1: torch.FloatTensor, shape = [1,2]
f2: torch.FloatTensor, shape = [1,2]
Returns
f3: torch.FloatTensor, shape = [1,2]
``
...

但是我的数据是批处理的，所以我有 f1 和 f2 的形状是 [8,1,2]

我想应用 my_func 得到形状为 [8,1,2] 的 f3。

这样做的惯用方法是什么？