class_docstr: ModuleBackend is a wrapper around torch module that allows it to run tasks asynchronously with Runtime
  By default, ModuleBackend handles three types of requests:
  
   - forward - receive inputs and compute outputs. Concurrent requests will be batched for better GPU utilization.
   - backward - receive gradients w.r.t. outputs, compute gradients w.r.t. inputs and **update expert**. Also batched.
   - get_info - return expert metadata. Not batched.
  
  :param module: nn.Module to be wrapped into a backend. Arbitrary pytorch module with a few limitations:
  
   - Experts must always receive the same set of args and kwargs and produce output tensors of same type
   - All args, kwargs and outputs must be **tensors** where 0-th dimension represents to batch size
   - We recommend using experts that are ~invariant to the order in which they process batches
   - Using randomness (e.g. Dropout) leads to different samples at forward and backward. If you want consistency,
      you should explicitly register these random variables as model inputs or outputs.
      See hivemind.utils.custom_layers.DeterministicDropout for an example
  
  :param optimizer: torch optimizer to be applied on every backward call
  :param scheduler: a function to create the learning rate scheduler for the expert
  :param args_schema: description of positional arguments to expert.forward, list of BatchTensorProto
  :param kwargs_schema: description of keyword arguments to expert.forward, dict of BatchTensorProto
  :param outputs_schema: description of outputs from expert.forward, nested structure of BatchTensorProto
  :param kwargs: extra parameters to be forwarded into TaskPool.__init__
class_docstr_tok: ["ModuleBackend","is","a","wrapper","around","torch","module","that","allows","it","to","run","tasks","asynchronously","with","Runtime","By","default",",","ModuleBackend","handles","three","types","of","requests",":","-","forward","-","receive","inputs","and","compute","outputs",".","Concurrent","requests","will","be","batched","for","better","GPU","utilization",".","-","backward","-","receive","gradients","w.r.t",".","outputs",",","compute","gradients","w.r.t",".","inputs","and","*","*","update","expert","*","*",".","Also","batched",".","-","get_info","-","return","expert","metadata",".","Not","batched",".",":","param","module",":","nn.Module","to","be","wrapped","into","a","backend",".","Arbitrary","pytorch","module","with","a","few","limitations",":","-","Experts","must","always","receive","the","same","set","of","args","and","kwargs","and","produce","output","tensors","of","same","type","-","All","args",",","kwargs","and","outputs","must","be","*","*","tensors","*","*","where","0-th","dimension","represents","to","batch","size","-","We","recommend","using","experts","that","are","~invariant","to","the","order","in","which","they","process","batches","-","Using","randomness","(","e.g",".","Dropout",")","leads","to","different","samples","at","forward","and","backward",".","If","you","want","consistency",",","you","should","explicitly","register","these","random","variables","as","model","inputs","or","outputs",".","See","hivemind.utils.custom_layers.DeterministicDropout","for","an","example",":","param","optimizer",":","torch","optimizer","to","be","applied","on","every","backward","call",":","param","scheduler",":","a","function","to","create","the","learning","rate","scheduler","for","the","expert",":","param","args_schema",":","description","of","positional","arguments","to","expert.forward",",","list","of","BatchTensorProto",":","param","kwargs_schema",":","description","of","keyword","arguments","to","expert.forward",",","dict","of","BatchTensorProto",":","param","outputs_schema",":","description","of","outputs","from","expert.forward",",","nested","structure","of","BatchTensorProto",":","param","kwargs",":","extra","parameters","to","be","forwarded","into","TaskPool.__init__"]
docstr: Return the current state of the module, optimizer, and scheduler
docstr_tok: ["Return","the","current","state","of","the","module",",","optimizer",",","and","scheduler"]
code: def state_dict(self) -> Dict:
      """Return the current state of the module, optimizer, and scheduler"""
      full_state = dict(module=self.module.state_dict())
      if self.optimizer is not None:
          full_state["optimizer"] = self.optimizer.state_dict()
      if self.scheduler is not None:
          full_state["scheduler"] = self.scheduler.state_dict()
      return full_state
  
file_path: hivemind/hivemind/moe/server/module_backend.py
imports: from typing import Any, Dict, Optional, Sequence, Tuple, Union
  import torch
  from torch import nn
  from hivemind.moe.server.task_pool import TaskPool
  from hivemind.utils.logging import get_logger
  from hivemind.utils.nested import nested_compare, nested_flatten, nested_map, nested_pack
  from hivemind.utils.tensor_descr import DUMMY_BATCH_SIZE, BatchTensorDescriptor
label_desc: Use image node_id 5 for calling the ModuleBackend obj's underlying member method code with example usage: obj.state_dict() and returns: full_state