code: def init_optim_state(optimizer: torch.optim.Optimizer) -> None:
      """
      Initialize optimizer states by calling step() with zero grads. This is necessary because some optimizers like AdamW
      initialize some states in their state_dicts lazily, only after calling step() for the first time. Certain checkpointing
      solutions may rely on in-place loading, re-using existing tensor allocated memory from the optimizer state dict. This
      optimization does not work with optimizers that lazily initialize their states, as certain states will not be restored.
      Calling this function ensures that these states are available in the state dict for in place loading.
  
      Args:
          optimizer: A PyTorch optimizer.
      """
      if optimizer.state:
          # The optimizer state is initialized.
          return
  
      for param_group in optimizer.param_groups:
          for param in param_group["params"]:
              if param.grad is not None:
                  raise RuntimeError(
                      "Initializing the optimizer states requires that no existing gradients for parameters are found."
                  )
              if param.requires_grad:
                  param.grad = torch.zeros_like(param)
      optimizer.step(closure=None)
      optimizer.zero_grad(set_to_none=True)
  
imports: from typing import Dict
  import torch
label_desc: Use image node_id 1 for calling a global function with example usage: init_optim_state(optimizer) without return types