SequenceModel step function _setup_state missing?

[FloMru]

I am building an encoder/decoder architecture where the encoder and decoder are S4 models. Since this is for a generative task I want to use the step function defined in the SequenceModel class for inference.

When I call the function I get: File "state-spaces/statespaces/models/sequence/model.py", line 129, in step x, state = layer.step(x, state=prev_state, **kwargs) File "state-spaces/src/models/sequence/block.py", line 117, in step y, state = self.layer.step(y, state, **kwargs) File "/state-spaces/src/models/sequence/ss/s4.py", line 253, in step y, next_state = self.kernel.step(u, state) # (B C H) File "/state-spaces/src/models/sequence/ss/kernel.py", line 1127, in step y, state = self.kernel.step(u, state, **kwargs) File "/state-spaces/src/models/sequence/ss/kernel.py", line 862, in step next_state = contract("h n, b h n -> b h n", self.dA, state)
File "/.local/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1269, in getattr raise AttributeError("'{}' object has no attribute '{}'".format( AttributeError: 'SSKernelDiag' object has no attribute 'dA'. Did you mean: '_A'?

Could the problem be, that the _setup_step function for the kernel is not called? If yes, what would be a practical way of calling the function, while using the SequenceModel class?

Thanks Flo