Mini batch size and reset states
You can simply make your training data a shape that is divisible by 50.
    In [1]: sample_size = 63648

In [2]: to_remove = sample_size % 50

In [3]: end_index = sample_size - to_remove

In [4]: model.fit(trainX[:end_index], trainY[:end_index], ...




What are the dimensions of trainX and trainY? You are not specifying a validation dataset, is that intentional?

You say you want to reset the state after a mini-batch, but you are resetting the state after each epoch (one loop through your entire dataset).

One more comment: the n_epochs = 10000 seems excessive.