Passing another parameter to LSTM


I am implementing a model from a paper, in which they modify the hidden state of LSTMCell by multiplying with a matrix d_c so that for a sequence, this would become the previous state for the next timestep h_{t+1}.The equation is:
h_t = [r_t ⊗ Φ(c_t)] ⊗ d_c

I understand can implement a custom LSTMCell to perform this task, but when I tried to make a custom LSTM, its speed is much slower than that implemented by Pytorch. Is there a way to still keep the speed of the Pytorch’s model while incorporating this new change?