parameters: (self, closure)
return_statement: return set(loaded), set(skipped)
docstring: r"""Advance state_dict load through callable ``closure`` whose signature is
          ``closure(key: str, var: Tensor) -> Union[np.ndarry, None]``
docstring_summary: r"""Advance state_dict load through callable ``closure`` whose signature is
          ``closure(key: str, var: Tensor) -> Union[np.ndarry, None]``
docstring_tokens: r Advance state_dict load through callable closure whose signature is closure ( key : str var : Tensor ) - > Union [ np . ndarry None ]
function: def _load_state_dict_with_closure(self, closure):
          r"""Advance state_dict load through callable ``closure`` whose signature is
          ``closure(key: str, var: Tensor) -> Union[np.ndarry, None]``
          """
          XNorm_typeclass = _get_XNorm_typeclass()
          assert callable(closure), "closure must be a function"
  
          loaded = []
          skipped = []
  
          local_state_dict = self._state_dict(keep_var=True)
          for (module_type, k), var in local_state_dict.items():
              to_be_load = closure(k, var)
              if to_be_load is None:
                  skipped.append(k)
                  continue
              assert isinstance(
                  to_be_load, np.ndarray
              ), "closure should return a `np.ndarray`, now `{}` get {}".format(
                  k, to_be_load
              )
              var_shape = make_shape_tuple(var.shape)
              to_be_load_shape = make_shape_tuple(to_be_load.shape)
              if var_shape != to_be_load_shape:
                  # weight and bias in BatchNorm1d, BatchNorm2d and SyncBatchNorm are 1-dim tensors in v1.0, and
                  # since v1.1 they are 4-dim tensors. The following special rule for these modules preserves the
                  # backward compatibility.
                  if issubclass(module_type, XNorm_typeclass):
                      if np.prod(var_shape) == np.prod(to_be_load_shape):
                          to_be_load = to_be_load.reshape(var_shape)
                      else:
                          raise ValueError(
                              "param `{}` size mismatch, should be {}, get {}".format(
                                  k, np.prod(var_shape), np.prod(to_be_load_shape)
                              )
                          )
                  else:
                      raise ValueError(
                          "param `{}` shape mismatch, should be {}, get {}".format(
                              k, var_shape, to_be_load_shape
                          )
                      )
              var._reset(
                  type(var)(
                      to_be_load, dtype=to_be_load.dtype, device=var.device, no_cache=True
                  )
              )
              loaded.append(k)
  
          return set(loaded), set(skipped)
function_tokens: def _load_state_dict_with_closure ( self , closure ) : XNorm_typeclass = _get_XNorm_typeclass ( ) assert callable ( closure ) , "closure must be a function" loaded = [ ] skipped = [ ] local_state_dict = self . _state_dict ( keep_var = True ) for ( module_type , k ) , var in local_state_dict . items ( ) : to_be_load = closure ( k , var ) if to_be_load is None : skipped . append ( k ) continue assert isinstance ( to_be_load , np . ndarray ) , "closure should return a `np.ndarray`, now `{}` get {}" . format ( k , to_be_load ) var_shape = make_shape_tuple ( var . shape ) to_be_load_shape = make_shape_tuple ( to_be_load . shape ) if var_shape != to_be_load_shape : # weight and bias in BatchNorm1d, BatchNorm2d and SyncBatchNorm are 1-dim tensors in v1.0, and # since v1.1 they are 4-dim tensors. The following special rule for these modules preserves the # backward compatibility. if issubclass ( module_type , XNorm_typeclass ) : if np . prod ( var_shape ) == np . prod ( to_be_load_shape ) : to_be_load = to_be_load . reshape ( var_shape ) else : raise ValueError ( "param `{}` size mismatch, should be {}, get {}" . format ( k , np . prod ( var_shape ) , np . prod ( to_be_load_shape ) ) ) else : raise ValueError ( "param `{}` shape mismatch, should be {}, get {}" . format ( k , var_shape , to_be_load_shape ) ) var . _reset ( type ( var ) ( to_be_load , dtype = to_be_load . dtype , device = var . device , no_cache = True ) ) loaded . append ( k ) return set ( loaded ) , set ( skipped )