wilson-labs/cola

device is not callable from Array

Closed this issue · 1 comments

hkortier commented

cola/cola/backends/jax_fns.py

Line 166 in 9562ae1

return array.device()

in jax 0.4.31 device is not callable.

AndPotap commented

Thanks for bringing this up. I just pushed a fix for this 😉