code: def step(
      self,
      closure: Optional[Callable[[], torch.Tensor]] = None,
      grad_clip_scale: Optional[Union[torch.Tensor, float]] = None,
  ):
      if not self._pointer_buffers_initialized:
          self._build_pointer_buffers()
  
      loss = closure() if closure is not None else None
  
      group = self.param_groups[0]
      compute_group = self.compute_param_groups[0]
      swa_group = self.swa_param_groups[0]
      if "step" in group:
          group["step"] += 1
      else:
          group["step"] = 1
      (beta1, beta2), step = group["betas"], group["step"]
      if group["bias_correction"]:
          beta1_correction = 1.0 - beta1**step
          beta2_correction = 1.0 - beta2**step
      else:
          beta1_correction = beta2_correction = 1.0
  
      grad_ptr = []
      for p in compute_group["params"]:
          if p.grad is None:
              continue
          if p.grad.detach().is_sparse:
              raise RuntimeError(
                  "FusedAdamSWA does not support sparse gradients, please consider SparseAdam instead"
              )
          grad_ptr.append(p.grad.data_ptr())
  
      for (
          compute_dtype,
          state_dtype,
      ), buffer_group in self.pointer_buffer_groups.items():
          device = buffer_group["device"]
          t_idx = buffer_group["tensor_idx"]
          grad_ptr_this_group = [grad_ptr[i] for i in t_idx]
          grad_ptr_this_group = torch.tensor(
              grad_ptr_this_group, dtype=torch.int64
          )
          grad_ptr_per_chunk = torch.repeat_interleave(
              grad_ptr_this_group, buffer_group["chunks_per_param"]
          ).to(device, non_blocking=True)
          if grad_clip_scale is None:
              grad_clip_scale_this_group = buffer_group[
                  "default_grad_clip_scale"
              ]
          elif not torch.is_tensor(grad_clip_scale):
              grad_clip_scale_this_group = torch.tensor(
                  grad_clip_scale
              ).to(device, non_blocking=True)
          else:
              grad_clip_scale_this_group = grad_clip_scale
  
          grid = (buffer_group["total_chunks"],)
          _multi_tensor_adam_swa[grid](
              state_param_ptr_per_chunk=buffer_group[
                  "param_ptr_per_chunk"
              ],
              compute_param_ptr_per_chunk=buffer_group[
                  "compute_param_ptr_per_chunk"
              ],
              swa_param_ptr_per_chunk=buffer_group[
                  "swa_param_ptr_per_chunk"
              ],
              grad_ptr_per_chunk=grad_ptr_per_chunk,
              moment_ptr_per_chunk=buffer_group[
                  "exp_avg_ptr_per_chunk"
              ],
              velocity_ptr_per_chunk=buffer_group[
                  "exp_avg_sq_ptr_per_chunk"
              ],
              chunk_local_idx_ptr=buffer_group["chunk_local_idx"],
              chunk_numel_ptr=buffer_group["chunk_numel"],
              grad_clip_scale_ptr=grad_clip_scale_this_group,
              lr=group["lr"],
              beta1=beta1,
              beta2=beta2,
              eps=group["eps"],
              weight_decay=group["weight_decay"],
              beta1_correction=beta1_correction,
              beta2_correction=beta2_correction,
              swa_decay_rate=self.swa_decay_rate,
              swa_n_averaged=swa_group["n_averaged"],
              adam_math_mode=self.adam_math_mode.value,
              MODEL_COMPUTE_DTYPE=_TORCH2DTYPE[compute_dtype],
              MODEL_STATE_DTYPE=_TORCH2DTYPE[state_dtype],
              # TODO: Find optimal hyper-parameters.
              CHUNK_SIZE=CHUNK_SIZE.item(),
              BLOCK_SIZE=128,
              num_warps=1,
          )
  
      swa_group["n_averaged"] += 1
  
      return loss
  
imports: from __future__ import annotations
  from collections import defaultdict
  from enum import Enum, unique
  from itertools import chain
  from typing import Callable, List, Optional, Tuple, Union
  import torch
  import torch.nn
  import triton
  import triton.language
  from torch.optim import Adam, Optimizer
label_desc: Use image node_id 3 for calling the FusedAdamSWA obj's underlying member method code with example usage: obj.step(closure, grad_clip_scale) and returns: loss