Hi Frank,
Thank you so much. That was really helpful. My goal is use this as an activation in building a neural network. Suppose I define a simple network like the following where I use nn.ReLU() activation. How do I package the idea of topk together so that I can create ‘mypackage’, import it and use mypackage.topk(k) instead of nn.ReLU()?
class Net(nn.Module):
    def __init__(self, input_size, hidden_size, out_size):
        super(Net, self).__init__()                    
        self.fc1 = nn.Linear(input_size, hidden_size)                              
        self.fc2 = nn.Linear(hidden_size, out_size)
        self.relu = nn.ReLU()

    def forward(self, x):                          
        out = self.fc1(x)
        out = self.relu(out)
        out = self.fc2(out)
        return out