How could I set different factor along one dimension using pow?

Hi,

I can use torch.pow like this:

factor = 2.3
a = torch.randn(2, 4, 2, 2)
b = a.pow(factor)

However, I need to use something like this:

factor = torch.tensor([2.1, 2.2, 2.3, 2.4])
a = torch.randn(2, 4, 2, 2)
b = a.pow(factor, dim=1)

It seems that pytorch does not support this. How could I do that along with its backward computation?

:grin:

b = a.pow(factor.broadcast_to((2,4)).unsqueeze(2).broadcast_to((2,4,2)).unsqueeze(3).broadcast_to((2,4,2,2)))

1 Like