Potential float mismatch between training inputs and model params

[kbcoulter]

Possible dtype mismatch. Model parameters are bfloat16 but training inputs are float32.

I believe that the .float() linked below casts to float32 and breaks anytime model params were created with --bfloat16.

embedding-kit/src/embkit/optimize/__init__.py

Line 322 in c41ffd5

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