code: def disable_apex_o2_state_dict_hook(
      model: Union[torch.nn.Module, torch.jit.ScriptFunction]
  ):
      # Apex O2 hook state_dict to return fp16 weights as fp32.
      # Exporter cannot identify them as same tensors.
      # Since this hook is only used by optimizer, it is safe to
      # remove this hook while exporting.
      if not isinstance(model, torch.jit.ScriptFunction):
          model_hooks = {}  # type: ignore[var-annotated]
          for module in model.modules():
              for key, hook in module._state_dict_hooks.items():
                  if type(hook).__name__ == "O2StateDictHook":
                      if module not in model_hooks:
                          model_hooks[module] = {}
                      model_hooks[module][key] = hook
              if module in model_hooks:
                  for key in model_hooks[module]:
                      module._state_dict_hooks.pop(key)
          try:
              yield
          finally:
              # Add the hooks back
              for module, m_map in model_hooks.items():
                  for key, hook in m_map.items():
                      module._state_dict_hooks[key] = hook
      else:
          try:
              yield
          finally:
              pass
  
imports: from __future__ import annotations
  import contextlib
  import copy
  import inspect
  import io
  import re
  import textwrap
  import typing
  import warnings
  from typing import Any, Callable, cast, Collection, Dict, List, Mapping, Optional, Sequence, Set, Tuple, Type, Union
  import torch
  import torch._C._onnx
  import torch.jit._trace
  import torch.serialization
  from torch import _C
  from torch.onnx import _constants, _exporter_states, errors, symbolic_caffe2, symbolic_helper
  from torch.onnx._globals import GLOBALS
  from torch.onnx._internal import _beartype, diagnostics, jit_utils, onnx_proto_utils, registration
label_desc: Use image node_id 3 for calling a global function with example usage: disable_apex_o2_state_dict_hook(model) without return types