S4 Module incompatible with pytorch 2.0's `torch.compile`

[RoiEXLab]

Hi,

I'm using the sashimi model on my own dataset with reasonable success for a while now and I wanted to see if I could use the recently released torch.compile function on the sashimi model to speed up training for my experiments.

Unfortunately it doesn't work. The following line seems to fail (for reasons I don't understand): https://github.com/HazyResearch/state-spaces/blob/06dbbdfd0876501a7f12bf3262121badbc7658af/src/models/s4/s4.py#L703

On the pytorch site there's some information on how to deal with those issues, so I hope the code can be extended in the future to run faster by a noticeable amount.

Thanks in advance.