Reproduction

Entropies is a bfloat16 tensor when training with bf16, but quantile() input tensor must be either float or double dtype.

System Info

RuntimeError: quantile() input tensor must be either float or double dtype

Checklist

• I have checked that my issue isn't already filed (see open issues)
• I have included my system information
• Any code provided is minimal, complete, and reproducible (more on MREs)
• Any code provided is properly formatted in code blocks, (no screenshot, more on code blocks)
• Any traceback provided is complete