code: def initialize_optimizer_state_(opt: torch.optim.Optimizer):
      """Initialize optimizer statistics by running a virtual optimizer step with zero gradients"""
      flat_params = tuple(
          param
          for group in opt.param_groups
          for param in group["params"]
      )
      old_grads = []
      for param in flat_params:
          old_grads.append(param.grad)
          param.grad = torch.zeros_like(param)
      opt.step()
      for param, old_grad in zip(flat_params, old_grads):
          param.grad = old_grad
  
imports: import logging
  import threading
  import time
  from concurrent.futures import ThreadPoolExecutor
  from contextlib import nullcontext
  from itertools import chain
  from typing import Any, Callable, Dict, Iterable, Iterator, Optional, Sequence, Tuple, Union
  import torch
  from packaging.version import Version
  import hivemind
  from hivemind.averaging import DecentralizedAverager
  from hivemind.averaging.control import StepControl
  from hivemind.compression import CompressionInfo, TensorRole
  from hivemind.optim.grad_scaler import GradScaler
  from hivemind.utils import DHTExpiration, PerformanceEMA, get_logger, nested_flatten, nested_pack
label_desc: Use image node_id 1 for calling a global function with example usage: initialize_optimizer_state_(opt) without return types