Tfa.layers.ESN Example

Does anyone have a working example implementing Each State Network with Tfa.layers.ESN.

Currently, I have the following model

Model: "model"
_________________________________________________________________
Layer (type)                 Output Shape              Param #   
_________________________________________________________________
digits (InputLayer)          [(None, 41, 4)]           0         
_________________________________________________________________
ESN                          (None, 41, 242)           30492     
_________________________________________________________________
flatten (Flatten)            (None, 9922)              0         
_________________________________________________________________
predictions (Dense)          (None, 2)                 19846  
_________________________________________________________________   
Total params: 50,338
Trainable params: 19,846
Non-trainable params: 30,492
_________________________________________________________________

The problem that I am having is that once I increase the number of hidden units the accuracy of the model drops instead of increasing, both in training and testing dataset. Also, the accuracy that I got is very low.