Python PyTorch-将不同功能应用于不同'；行/列'；张量

假设我有一个二维张量：

x = torch.Tensor([[1, 2], [3, 4]])

是否有一种有效的方法将一个函数应用于第一行[1，2]，并将第二个不同的函数应用于第二行[3，4]？（不必是一行，可以跨越任何维度）

目前，我使用以下代码：假设我有两个函数，例如f和g

def f(z):
    return 2 * z
def g(z):
    return 0.5 * z

然后，要将它们应用于单独的行，我将执行以下操作：

torch.cat([f(x[0]).unsqueeze(0), g(x[1]).unsqueeze(0)], dim = 0)

它给出了所需的张量[[2,4]，[1.5,2]]

---

显然，在这个二维示例中，这个解决方案很好，但似乎有点笨拙。有更好的方法吗？特别是在更高的维度中，或者当所选维度中有大量元素时

如果问题只是因为太笨重，您可以将拆分和串联步骤推到函数内部，而不是在函数外部执行。