code: def _single_tensor_sophiag(
      params: List[Tensor],
      grads: List[Tensor],
      exp_avgs: List[Tensor],
      hessian: List[Tensor],
      state_steps: List[Tensor],
      *,
      bs: int,
      beta1: float,
      beta2: float,
      rho: float,
      lr: float,
      weight_decay: float,
      maximize: bool,
      capturable: bool,
  ):
      for i, param in enumerate(params):
          grad = grads[i] if not maximize else -grads[i]
          exp_avg = exp_avgs[i]
          hess = hessian[i]
          step_t = state_steps[i]
  
          if capturable:
              assert param.is_cuda and step_t.is_cuda and bs.is_cuda
  
          if torch.is_complex(param):
              grad = torch.view_as_real(grad)
              exp_avg = torch.view_as_real(exp_avg)
              hess = torch.view_as_real(hess)
              param = torch.view_as_real(param)
  
          # update step
          step_t += 1
  
          # Perform stepweight decay
          param.mul_(1 - lr * weight_decay)
  
          # Decay the first and second moment running average coefficient
          exp_avg.mul_(beta1).add_(grad, alpha=1 - beta1)
  
          if capturable:
              step = step_t
              step_size = lr
              step_size_neg = step_size.neg()
  
              ratio = (exp_avg.abs() / (rho * bs * hess + 1e-15)).clamp(
                  None, 1
              )
              param.addcmul_(exp_avg.sign(), ratio, value=step_size_neg)
          else:
              step = step_t.item()
              step_size_neg = -lr
  
              ratio = (exp_avg.abs() / (rho * bs * hess + 1e-15)).clamp(
                  None, 1
              )
              param.addcmul_(exp_avg.sign(), ratio, value=step_size_neg)
  
file_path: bert-ultrafast/training/cramming/backend/optimizers/sophia.py
imports: import torch
  from torch import Tensor
  from torch.optim.optimizer import Optimizer
  from typing import List
label_desc: Use image node_id 2 for calling a global function with example usage: _single_tensor_sophiag(params, grads, exp_avgs, hessian, state_steps) without return types