docstr: Returns the state of SWA as a :class:`dict`.
  
  It contains three entries:
      * opt_state - a dict holding current optimization state of the base
          optimizer. Its content differs between optimizer classes.
      * swa_state - a dict containing current state of SWA. For each
          optimized variable it contains swa_buffer keeping the running
          average of the variable
      * param_groups - a dict containing all parameter groups
docstr_tok: ["Returns","the","state","of","SWA","as","a",":","class",":","`","dict","`",".","It","contains","three","entries",":","*","opt_state","-","a","dict","holding","current","optimization","state","of","the","base","optimizer",".","Its","content","differs","between","optimizer","classes",".","*","swa_state","-","a","dict","containing","current","state","of","SWA",".","For","each","optimized","variable","it","contains","swa_buffer","keeping","the","running","average","of","the","variable","*","param_groups","-","a","dict","containing","all","parameter","groups"]
code: def state_dict(self):
      r"""Returns the state of SWA as a :class:`dict`.
  
      It contains three entries:
          * opt_state - a dict holding current optimization state of the base
              optimizer. Its content differs between optimizer classes.
          * swa_state - a dict containing current state of SWA. For each
              optimized variable it contains swa_buffer keeping the running
              average of the variable
          * param_groups - a dict containing all parameter groups
      """
      opt_state_dict = self.optimizer.state_dict()
      swa_state = {
          (id(k) if isinstance(k, torch.Tensor) else k): v
          for k, v in self.state.items()
      }
      opt_state = opt_state_dict["state"]
      param_groups = opt_state_dict["param_groups"]
      return {
          "opt_state": opt_state,
          "swa_state": swa_state,
          "param_groups": param_groups,
      }
  
imports: import warnings
  from collections import defaultdict
  from typing import Optional, Union
  import torch
  from pytext.config.component import create_optimizer
  from pytext.optimizer.lamb import Lamb
  from pytext.optimizer.madgrad import MADGRAD
  from pytext.optimizer.optimizers import Adagrad, Adam, AdamW, Optimizer, SGD
  from pytext.optimizer.radam import RAdam
  from torch.optim import Optimizer
label_desc: Use image node_id 8 for calling the StochasticWeightAveraging obj's underlying member method code with example usage: obj.state_dict() and returns: dict