function for Ordinal Pooling Neural network

please I want to create a function that computes the Ordinal Pooling neural network like the following figure:

this is my function :

def Ordinal_Pooling_NN(x):
  wights = torch.tensor([0.6, 0.25, 0.10, 0.05])
  top = torch.topk(x, 4, dim = 1)
  wights = wights.repeat(x.shape[0], 1)
  result = torch.sum(wights * (top.values), dim = 1 )
  return result

but as a result, I get the following error:

<ipython-input-112-ddf99c812d56> in Ordinal_Pooling_NN(x)
      9     top = torch.topk(x, 4, dim = 1)
     10     wights = wights.repeat(x.shape[0], 1)
---> 11     result = torch.sum(wights * (top.values), dim = 1 )
     12     return result

RuntimeError: The size of tensor a (4) must match the size of tensor b (16) at non-singleton dimension 2

Your implementation is actually correct, I believe you did not feed the function with a 2D tensor, the input must have a batch axis. For instance, the code below will run:

>>> Ordinal_Pooling_NN(torch.tensor([[1.9, 0.4, 1.3, 0.8]]))
tensor([1.5650])

Do note you are not required to repeat the weights tensor, it will be broadcasted automatically when computing the point-wise multiplication. You only need the following:

def Ordinal_Pooling_NN(x):
  w = torch.tensor([0.6, 0.25, 0.10, 0.05])
  top = torch.topk(x, k=4, dim=1)
  result = torch.sum(w*top.values, dim=1)
  return result