Using rnn output in a custom layer?

I am trying to make my own customer layer with an RNN that takes the generated text, does some transformations to it, then converts back into a tensor. I am not sure if what Im thinking is doable. Is it possible to generate the output of the current batch inside an RNN model during a forward pass?
I want to use that data to inform the loss function.
I tried with this model definition:

class TestLayer(nn.Module):
    def __init__(self):
    def forward(self,input):
       #hidden = decoder.init_hidden()
        prime_input = char_tensor(prime_str)
        predicted = prime_str

        # Use priming string to "build up" hidden state
        for p in range(len(prime_str) - 1):
            _, hidden = decoder(prime_input[p], hidden)

        inp = prime_input[-1]

        for p in range(predict_len):
            output, hidden = decoder(inp, hidden)

            # Sample from the network as a multinomial distribution
            output_dist =
            top_i = torch.multinomial(output_dist, 1)[0]

            # Add predicted character to string and use as next input
            predicted_char = all_characters[top_i]
           predicted += predicted_char
           inp = char_tensor(predicted_char)

        return predicted
class RNN(nn.Module):
    def __init__(self, input_size, hidden_size, output_size, n_layers=1):
        super(RNN, self).__init__()
        self.input_size = input_size
        self.hidden_size = hidden_size
        self.output_size = output_size
        self.n_layers = n_layers
        self.test_layer = TestLayer()

        self.encoder = nn.Embedding(input_size, hidden_size)
        self.gru = nn.GRU(hidden_size, hidden_size, n_layers)
        self.decoder = nn.Linear(hidden_size, output_size)

    def forward(self, input, hidden):
        input = self.encoder(input.view(1, -1))
        output, hidden = self.gru(input.view(1, 1, -1), hidden)
        output = self.decoder(output.view(1, -1))
        output = self.test_layer(output)
        return output, hidden

    def init_hidden(self):
        return Variable(torch.zeros(self.n_layers, 1, self.hidden_size))