I want to implement a custom activation function with learnable parameters. For example one that takes the input x and returns a polinomial of specified order, of x.
def poli_activation(x, order=2):
  input_tens = []

  # is this the way to make coeff a vector of parameters?
  coeff = torch.nn.parameter(torch.randn(order+1))

  # need a vector of powers of x , for example (x^2, x, 1)
  for idx in range(order):
    element = torch.pow(x, idx)
    input_tens.append(element)

  # make this vector a variable and implement polynomial
  # as dot product.
  input_tens = torch.Tensor(input_tens)
  input_tens = torch.autograd.Variable(input_tens.cuda())
  output = torch.dot(coeff, input_tens)
  return output

Now this is 1) a very inefficient implementation of such a function and 2) does not work. Does anyone have an idea of how this can be efficiently implemented in pytorch? I want the network to be able to learn the polynomial coefficients.