Generating Names with a Character-Level RNN modification

I am fairly new to using RNNs and pytorch. I study the Generating Names with a Character-Level RNN tutorial and I was wondering how should I change the model to use an LSTMCell. The model in the tutorial is defined as follows:

class RNN(nn.Module):
    def __init__(self, input_size, hidden_size, output_size):
        super(RNN, self).__init__()
        self.hidden_size = hidden_size

        self.i2h = nn.Linear(n_categories + input_size + hidden_size, hidden_size)
        self.i2o = nn.Linear(n_categories + input_size + hidden_size, output_size)
        self.o2o = nn.Linear(hidden_size + output_size, output_size)
        self.dropout = nn.Dropout(0.1)
        self.softmax = nn.LogSoftmax(dim=1)

I am thinking to replace self.i2h with the LSTMCell layer:

self.lstm = nn.LSTMCell(n_categories + input_size + hidden_size, hidden_size)

I am not sure then if self.i2o should be kept or become something like that:

self.h2o = nn.Linear(hidden_size, output_size)