Support alternative precision training (#333)

Update matrix to use the dtype of the gradient during projection to fix error. Internal representation of float32 is maintained for precision purposes.

Author: kylevedder

pytorch_optimizer/optimizer/soap.py

@@ -102,7 +102,7 @@ def project(
 
         for mat in state['Q']:
             if len(mat) > 0:
-                grad = torch.tensordot(grad, mat, dims=[[0], [0 if project_type == 'forward' else 1]])
+                grad = torch.tensordot(grad, mat.to(grad.dtype), dims=[[0], [0 if project_type == 'forward' else 1]])
             else:
                 grad = grad.permute([*list(range(1, len(grad.shape))), 0])