johnma2006 / mamba-minimal

Simple, minimal implementation of the Mamba SSM in one file of PyTorch.
Apache License 2.0
2.62k stars 191 forks source link

Enable using the model on GPU #3

Closed gcaillaut closed 10 months ago

gcaillaut commented 10 months ago

This small change put x on the same device as the model parameters.