S4 Listops have nan loss

[lthilnklover]

First of all, thank you for the comprehensive code base for all variants of S4 models.

However, as I try to run the Listops experiments with S4 (HYYT version), the losses for train, test and val all become nan after 1 epoch.

I ran the following script:

python -m train experiment=lra/s4-listops wandb=null

The final accuracy is also way below the reported accuracy (train=0.17).

Is there something that I have done wrong..?