docstr: For given masks, update the masks in the wrappers.
  
  Parameters
  ----------
  masks
      The masks' format should be {module_name: {target_name: target_masks}}.
      If the module is not registered as a wrapper in the model or the target is not registered as a target space in the compressor,
      the related masks updating will be ignored.
docstr_tok: ["For","given","masks",",","update","the","masks","in","the","wrappers",".","Parameters","--","--","--","--","--","masks","The","masks","'","format","should","be","{","module_name",":","{","target_name",":","target_masks","}","}",".","If","the","module","is","not","registered","as","a","wrapper","in","the","model","or","the","target","is","not","registered","as","a","target","space","in","the","compressor",",","the","related","masks","updating","will","be","ignored","."]
code: def update_masks(self, masks: Dict[str, Dict[str, torch.Tensor]]):
      """
      For given masks, update the masks in the wrappers.
  
      Parameters
      ----------
      masks
          The masks' format should be {module_name: {target_name: target_masks}}.
          If the module is not registered as a wrapper in the model or the target is not registered as a target space in the compressor,
          the related masks updating will be ignored.
      """
      for module_name, target_masks in masks.items():
          if module_name not in self._module_wrappers:
              _logger.warning(
                  f"Wrapper {module_name} is not register in this compressor, the masks will be ignored."
              )
          wrapper = self._module_wrappers[module_name]
          for target_name, target_mask in target_masks.items():
              target_space = wrapper.pruning_target_spaces.get(
                  target_name, None
              )
              if target_space is None:
                  _logger.warning(
                      f"Target space `{target_name}` is not in wrapper `{wrapper.name}`, the mask of it will be ignored."
                  )
                  continue
              if target_mask is None:
                  target_space.mask = None
              else:
                  assert (
                      target_name in wrapper.pruning_target_spaces
                  ), f"{module_name}.{target_name} is not a pruning target, can not update mask for it."
                  # NOTE: try to auto get the device of the current module
                  try:
                      device = next(wrapper.parameters()).device
                  except StopIteration:
                      try:
                          device = next(wrapper.buffers()).device
                      except StopIteration:
                          if target_space.mask is not None:
                              device = target_space.mask.device
                          else:
                              # NOTE: this will have risk in model parallel
                              device = next(
                                  self.bound_model.parameters()
                              ).device
                  target_space.mask = target_mask.to(device)
  
imports: from __future__ import annotations
  from collections import defaultdict
  from copy import deepcopy
  import logging
  from typing import Any, Dict, List, Literal, Callable
  import torch
  from torch.utils._pytree import tree_map
  from .config import trans_legacy_config_list, default_config_schema
  from .target_space import TargetSpace, TargetType, DistillationTargetSpace, PruningTargetSpace, QuantizationTargetSpace
  from .wrapper import ModuleWrapper, register_wrappers
  from ..utils import Evaluator, Scaling, _EVALUATOR_DOCSTRING
label_desc: Use image node_id 5 for calling the Pruner obj's underlying member method code with example usage: obj.update_masks(masks) without return types