Pytorch version PI0.5 model mismatchs the jax converted pytorch checkpoint.

[XIEdeda]

Thank you for the open-source pi0.5 model and its PyTorch version. However, when attempting to perform inference using the PyTorch version of the pi0.5 model and its checkpoint, I encountered the following issues:I saw others say that it can be successfully converted in a uv environment, so why doesn't it work under conda?

RuntimeError: Error(s) in loading state_dict for PI0Pytorch: Invalid key(s) in state_dict: "action_in_proj.bias", "action_in_proj.weight", "action_out_proj.bias", "action_out_proj.weight", "paligemma_with_expert.gemma_expert.model.layers.0.input_layernorm.dense.bias", "paligemma_with_expert.gemma_expert.model.layers.0.input_layernorm.dense.weight", "paligemma_with_expert.gemma_expert.model.layers.0.post_attention_layernorm.dense.bias", "paligemma_with_expert.gemma_expert.model.layers.0.post_attention_layernorm.dense.weight", ...... "paligemma_with_expert.paligemma.model.vision_tower.vision_model.embeddings.patch_embedding.weight", "paligemma_with_expert.paligemma.model.vision_tower.vision_model.embeddings.position_embedding.weight", "time_mlp_in.bias", "time_mlp_in.weight", "time_mlp_out.bias", "time_mlp_out.weight", mismatched dtypes or shape.