PyTorch, применять различные функции поэлементно

Question

Я определил тензор как этот

t_shape = [4, 1]
data = torch.rand(t_shape)

Я хочу применить различные функции к каждой строке.

funcs = [lambda x: x+1, lambda x: x**2, lambda x: x-1, lambda x: x*2]  # each function for each row.

Я могу сделать это с помощью следующего кода

d = torch.tensor([f(data[i]) for i, f in enumerate(funcs)])

Как правильно сделать это с более продвинутыми API, определенными в PyTorch?

Answer 1

Я думаю, что ваше решение хорошо. Но это не будет работать с любой тензорной формой. Вы можете немного изменить решение следующим образом.

t_shape = [4, 10, 10]
data = torch.rand(t_shape)

funcs = [lambda x: x+1, lambda x: x**2, lambda x: x-1, lambda x: x*2]

# only change the following 2 lines
d = [f(data[i]) for i, f in enumerate(funcs)]
d = torch.stack(d, dim=0)