Merge branch 'main' into wan_transformer

Author: jfacevedo-google

src/maxdiffusion/models/modeling_flax_pytorch_utils.py

@@ -66,7 +66,8 @@ def torch2jax(torch_tensor: torch.Tensor) -> Array:
     torch_tensor = torch_tensor.to("cpu")
 
   numpy_value = torch_tensor.numpy()
-  jax_array = jnp.array(numpy_value, dtype=jnp.bfloat16 if is_bfloat16 else None)
+  local_cpu_device_0 = jax.local_devices(backend="cpu")[0]
+  jax_array = jnp.array(numpy_value, dtype=jnp.bfloat16 if is_bfloat16 else None, device=local_cpu_device_0)
   return jax_array