How to move train_state object to CPU?

[mattiasmar]

Building on FLAX' mnist example the train_state object is created by:
model = train.create_train_state(init_rng, config)

This allocates memory on the GPU. I tried moving this memory buffer to CPU by calling:

cpus = jax.devices("cpu")
jax.device_put(model.params, cpus[0])

and also by calling:
model.params = jax.tree_map(lambda x: np.asarray(x), model.params)

The former didn't release the memory on the GPU and the later gave FrozenInstanceError.
Any suggestions on how to move the memory allocation from the GPU to the CPU, or how to do the allocation directly on the CPU to begin with (without limiting the rest of the program to use the GPU)?

[jheek]

The easiest way to do this is to simply create the train_state on CPU to begin with. In this case by using model = jax.jit(train.create_train_state, static_argnums=(1,), backend="cpu")(init_rng, config)

Btw JAX is currently missing a way to control device placement without using jit see jax-ml/jax#8879

[mattiasmar]

Thanks!