fix device selection for mps

[xlr8harder]

This was causing the kvcache to end up on cpu on mps.

This doesn't entirely fix mac: the pyproject.toml tries to force a cuda version of torch, but I don't know the pyproject format well enough to know how to fix that cleanly.

[xlr8harder]

Ah I see there is already a pending PR with this fix: #41

[artus-LYTiQ]

o7 xlr8!

Try the following. toml line 18:

torch = [ {version = "^2.4.1", source = "pypi", markers = "platform_system == 'Darwin'"}, {version = "^2.4.1", source = "pytorch", extras = ["+cu124"], markers = "platform_system != 'Darwin'"} ]

Then add before or behind the other .source:

[[tool.poetry.source]] name = "pypi" priority = "primary"

with the respective line breaks of course. HTH

[xlr8harder]

oh, thanks! and that won't break the cuda installs separately?

[xlr8harder]

You should probably submit that fix as a separate PR!

[artus-LYTiQ]

I honestly still have to test it on ubuntu+4090 but I am still stuck in the jax-metal PR.