Lightning precision in orthogonalize()

[eric-tc-wong]

Hello,

Thank you for the great work.
I encountered an error when using Dion with pytorch lightning with precision="bf16-mixed".

dion/dion.py", line 1302, in orthogonalize

[rank1]:     R, _ = torch.linalg.cholesky_ex(PP, upper=True)

[rank1]:            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^

[rank1]: NotImplementedError: "cholesky_cuda" not implemented for 'BFloat16'

Quick fix that appears to work:

@torch.amp.autocast("cuda",enabled=False)
def orthogonalize(P: Tensor, oversample: float = 1.25) -> Tensor:
...

Thanks