Prevent PI0/PI05 JAX to PyTorch weight precision drift.

[Greyman-Seu]

Summary

This PR adds a small safeguard to the JAX-to-PyTorch checkpoint conversion path for PI0-family models, including PI0 and PI05.

During conversion, the model config dtype is forced to float32 internally if the source config uses another dtype. This avoids accidental precision loss in the intermediate PyTorch checkpoint and helps keep JAX and PyTorch weights
numerically aligned.

Motivation

During OpenPI inference acceleration work, we need the following path to stay numerically aligned:

JAX checkpoint -> PyTorch checkpoint -> downstream accelerated framework

For this pipeline, JAX and PyTorch should be able to produce fully aligned inference results when inputs, weights, and precision settings match. If the PyTorch checkpoint is built with a lower-precision model config during
conversion, fp32 weight precision can be lost before downstream conversion or debugging, making it harder to distinguish real backend differences from conversion artifacts.

This change adds a defensive guard so PI0/PI05 conversion restores and constructs the PyTorch model in fp32 internally, while keeping the existing CLI output precision behavior unchanged.

## Changes

- Import dataclasses
- During PI0-family checkpoint conversion, replace model_config.dtype with "float32" when needed
- Keep the existing precision CLI default and options unchanged

## Notes

This does not change the requested output precision. It only prevents the model config dtype from silently reducing internal conversion precision before the final save step.