When converting a slimmed model to the huggingface format with the CPU, I get the following assertion error:
File "to_hf_weights.py", line 488, in <module>
save_sharded_to_hf_format(input_ckpt, params, output_path, np_dtype, torch_dtype)
File "to_hf_weights.py", line 466, in save_sharded_to_hf_format
save_pytree_as_hf(
File "to_hf_weights.py", line 382, in save_pytree_as_hf
x = unshard_leave(x, leave_name, old_shape, np_dtype=np_dtype)
File "to_hf_weights.py", line 312, in unshard_leave
assert isinstance(x, jnp.ndarray)
AssertionError
When converting a slimmed model to the huggingface format with the CPU, I get the following assertion error:
Simply commenting out the line
works fine.