oneflow.nn.Graph.load_state_dict

Graph.load_state_dict(state_dict: Dict[str, Union[Dict[str, oneflow.Tensor], oneflow.Tensor]], strict: bool = True)

Copies module’s states and other graph states from state_dict into this graph. If strict is True, then the keys of state_dict must exactly match the keys returned by this module’s nn.Graph.state_dict() function.

Parameters
  • state_dict (dict) – a dict containing module’s states and other graph states.

  • strict (bool, optional) – whether to strictly enforce that the keys in state_dict match the keys returned by this graph’s nn.Graph.state_dict() function. Default: True.

Note

nn.Graph’s state dict can only be loaded before the first call of a graph.