refactor: pre_conditioners_for_grad type

Author: kozistr

pytorch_optimizer/optimizer/shampoo_utils.py

@@ -333,7 +333,7 @@ def compute_pre_conditioners(self) -> None:
     def precondition_block(
         partitioned_grad: torch.Tensor,
         should_preconditioned_dims: List[bool],
-        pre_conditioners_for_grad: List[torch.Tensor],
+        pre_conditioners_for_grad: Union[List[torch.Tensor], torch.Tensor],
     ) -> torch.Tensor:
         r"""Perform a preconditioning operation on a single gradient block.