class_docstr: Sharpness-Aware Minimization that simultaneously minimizes loss value and loss sharpness.
  In particular, it seeks parameters that lie in neighborhoods having uniformly low loss.
code: def first_step(self, zero_grad=False):
      grad_norm = self._grad_norm()
      for group in self.param_groups:
          scale = group["rho"] / (grad_norm + 1e-12)
  
          for p in group["params"]:
              if p.grad is None:
                  continue
              self.state[p]["old_p"] = p.data.clone()
              e_w = (
                  (torch.pow(p, 2) if group["adaptive"] else 1.0)
                  * p.grad
                  * scale.to(p)
              )
              p.add_(e_w)  # climb to the local maximum "w + e(w)"
  
      if zero_grad:
          self.zero_grad()
  
file_path: towhee/towhee/models/layers/sam.py
imports: import torch
label_desc: Use image node_id 2 for calling the SAM obj's underlying member method code with example usage: obj.first_step(zero_grad) without return types