@fmassa Can you please give an example on that?
I tried to implement a simple exponential activation for my rnn
class my_rnn(nn.Module):
    def __init__(self, input_size=2, hidden_size=20, num_layers=3, output_size=1,
                 batch_size=10):
        super(my_rnn, self).__init__()
        self.input_size = input_size
        self.hidden_size = hidden_size
        self.num_layers = num_layers
        self.output_size = output_size
        self.batch_size = batch_size
        
        self.rnn = nn.RNN(input_size=self.input_size, hidden_size=self.hidden_size, 
                          num_layers=self.num_layers, batch_first=True, nonlinearity=self.exp_activation)
        # The last layer is applied on the last output only of the RNN (not like
        # TimeDistributedDense in Keras)
        self.linear_layer = nn.Linear(self.hidden_size, self.output_size)
        
        self.hidden = Variable(torch.zeros((self.num_layers, self.batch_size, self.hidden_size))).cuda()
        
    def forward(self, input_sequence):
        out_rnn, self.hidden = self.rnn(input_sequence, self.hidden)
        in_linear = out_rnn[:, -1, :]
        final_output = self.linear_layer(in_linear)
        return final_output
    
    def init_hidden(self):
        self.hidden = Variable(torch.zeros((self.num_layers, self.batch_size, self.hidden_size))).cuda()

    def exp_activation(self, data):
        return torch.exp(data)

but it gives this error ValueError: Unknown nonlinearity '<bound method my_rnn.exp_activation of my_rnn ()>'