Maximilien Le Clei

it's just about picking which model to import from https://github.com/state-spaces/mamba/tree/main/mamba_ssm/modules however if you're using the MixerModel class, add `ssm_cfg={"layer": "Mamba2"}` as an argument

https://github.com/Lightning-AI/pytorch-lightning/blob/3740546899aedad77c80db6b57f194e68c455e28/src/lightning/pytorch/trainer/trainer.py#L146 I guess this section could add that info for CPU. So you mean `devices=1` on CPU equates to using all CPUs?

Hey both, thanks for the quick replies, will provide more info in the coming day(s), but days might turn into week(s) due to current time constraints :sweat_smile:

Soooooo, I tried reproducing the error a few times in the same environment but it was never raised again :ghost: . `jaxtyping==0.2.24`, `beartype==0.16.4` Anyways here is a type hint that...

Thanks for the detailed answer @leycec. Unfortunately this time I really won't have the bandwidth to test and report in the coming days. I'll edit my answer when I do...

hey @tridao, any chance you could build the `torch 2.7` wheels for `causal_conv1d` and `mamba_ssm`? and if you have access to an AMD GPU, ROCm wheels would also be welcomed...