parameters: (self, state_dict, pytorch_format)
docstring: Extract model states from the training session and load into the state_dict
docstring_summary: Extract model states from the training session and load into the state_dict
docstring_tokens: Extract model states from the training session and load into the state_dict
function: def _extract_model_states(self, state_dict, pytorch_format):
          """Extract model states from the training session and load into the state_dict"""
  
          model_states = self._training_session.get_model_state(include_mixed_precision_weights=False)
          state_dict[_utils.state_dict_model_key()] = {}
  
          # extract trained model weights from the training session
          for precision in model_states:
              state_dict[_utils.state_dict_model_key()][precision] = {}
              for model_state_key in model_states[precision]:
                  if pytorch_format:
                      state_dict[_utils.state_dict_model_key()][precision][model_state_key] = \
                          torch.from_numpy(model_states[precision][model_state_key])
                  else:
                      state_dict[_utils.state_dict_model_key()][precision][model_state_key] = \
                          model_states[precision][model_state_key]
  
          # extract untrained (frozen) model weights
          for node in self._onnx_model.graph.initializer:
              if node.name not in state_dict[_utils.state_dict_model_key()][_utils.state_dict_full_precision_key()] and \
                  node.name in self.options.utils.frozen_weights:
                  if pytorch_format:
                      state_dict[_utils.state_dict_model_key()][_utils.state_dict_full_precision_key()][node.name] = \
                          torch.from_numpy(onnx.numpy_helper.to_array(node))
                  else:
                      state_dict[_utils.state_dict_model_key()][_utils.state_dict_full_precision_key()][node.name] = \
                          onnx.numpy_helper.to_array(node)
function_tokens: def _extract_model_states ( self , state_dict , pytorch_format ) : model_states = self . _training_session . get_model_state ( include_mixed_precision_weights = False ) state_dict [ _utils . state_dict_model_key ( ) ] = { } # extract trained model weights from the training session for precision in model_states : state_dict [ _utils . state_dict_model_key ( ) ] [ precision ] = { } for model_state_key in model_states [ precision ] : if pytorch_format : state_dict [ _utils . state_dict_model_key ( ) ] [ precision ] [ model_state_key ] = torch . from_numpy ( model_states [ precision ] [ model_state_key ] ) else : state_dict [ _utils . state_dict_model_key ( ) ] [ precision ] [ model_state_key ] = model_states [ precision ] [ model_state_key ] # extract untrained (frozen) model weights for node in self . _onnx_model . graph . initializer : if node . name not in state_dict [ _utils . state_dict_model_key ( ) ] [ _utils . state_dict_full_precision_key ( ) ] and node . name in self . options . utils . frozen_weights : if pytorch_format : state_dict [ _utils . state_dict_model_key ( ) ] [ _utils . state_dict_full_precision_key ( ) ] [ node . name ] = torch . from_numpy ( onnx . numpy_helper . to_array ( node ) ) else : state_dict [ _utils . state_dict_model_key ( ) ] [ _utils . state_dict_full_precision_key ( ) ] [ node . name ] = onnx . numpy_helper . to_array ( node )