big_vision icon indicating copy to clipboard operation
big_vision copied to clipboard

Error with putting arrays on CPU in cloud TPUs

Open philippe-eecs opened this issue 9 months ago • 1 comments

Hi I've been setting up big vision on a v4-32 TPU pod and I run into this error whenever I call u.put_cpu

jaxlib.xla_extension.XlaRuntimeError: INVALID_ARGUMENT: Cannot copy array to non-addressable device TFRT_CPU_0

I'm guessing the CPUs on the TPU pod aren't configured properly? Is there a way around this or a way to fix this issue?

Totally new to TPUs and let me know if you need more information.

philippe-eecs avatar Apr 29 '24 21:04 philippe-eecs