Align fit_vae input casting with model parameter dtype

[Copilot]

fit_vae was force-casting training batches to float32, which can conflict with models initialized in bfloat16 (e.g., via CLI --bfloat16). This change removes the hardcoded cast and ensures batch tensors follow the model’s parameter dtype at training time.

• Training path: remove hardcoded fp32 cast

  • In fit_vae, replace x_tensor.to(device).float() with dtype-aware casting derived from model parameters.
  • Preserve device transfer behavior; only dtype policy changes.

• Dtype policy in fit_vae

  • Determine model_dtype once after moving model to device.
  • In each VAE step:
    • if model has parameters: cast batch to device + model_dtype
    • fallback for parameterless models: cast to device only

• Regression coverage

  • Added a focused optimize test asserting that float32 loader batches are converted to bfloat16 when the model parameters are bfloat16.

# before
x_tensor = x_tensor.to(device).float()

# after
if model_dtype is None:
    x_tensor = x_tensor.to(device)
else:
    x_tensor = x_tensor.to(device=device, dtype=model_dtype)

• Fixes Potential float mismatch between training inputs and model params #69

[github-actions]

☂️ Python Coverage

current status: ✅

Overall Coverage

Lines	Covered	Coverage	Threshold	Status
3192	2738	86%	0%	🟢

New Files

No new covered files...

Modified Files

File	Coverage	Status
src/embkit/optimize/_init_.py	82%	🟢
TOTAL	82%	🟢

updated for commit: 7732e92 by action🐍