I would suggest that you build your own LSTM cell for such experiments.
You could use the cell we use in benchmarking as a start 62 and then build your own nn.Module based on the function.
Best regards
Thomas