code: def _update_weights(self, res):
      if self.model_dir is not None:
          result = res
          try:
              temp_dir = tempfile.mkdtemp()
              get_remote_file_to_local(
                  os.path.join(self.model_dir, "state.pkl"),
                  os.path.join(temp_dir, "state.pkl"),
              )
              from bigdl.orca.common import SafePickle
  
              with open(os.path.join(temp_dir, "state.pkl"), "rb") as f:
                  state = SafePickle.load(f)
                  self.model_weights = state["weights"]
          finally:
              shutil.rmtree(temp_dir)
      else:
          result = res[0]
          states = res[1]
          self.model_weights = states["weights"]
          self.optimizer_weights = states["opt_weights"]
      return result
  
imports: import os
  import types
  import logging
  import tempfile
  import shutil
  from bigdl.dllib.utils.common import get_node_and_core_number
  from bigdl.dllib.utils.file_utils import is_local_path
  from bigdl.dllib.utils.utils import get_node_ip
  from bigdl.orca.data.file import is_file, exists, get_remote_file_to_local, get_remote_files_with_prefix_to_local, put_local_file_to_remote, put_local_files_with_prefix_to_remote
  from bigdl.orca.learn.tf2.spark_runner import SparkRunner
  from bigdl.orca.learn.utils import find_free_port, find_ip_and_free_port
  from bigdl.orca.learn.utils import maybe_dataframe_to_xshards, dataframe_to_xshards, convert_predict_xshards_to_dataframe, make_data_creator, load_model, save_model, process_xshards_of_pandas_dataframe, add_predict_to_pd_xshards, update_predict_xshards
  from bigdl.orca.learn.log_monitor import start_log_server, stop_log_server
  from bigdl.orca import OrcaContext
  from bigdl.orca.data.shard import SparkXShards
  from bigdl.dllib.utils.log4Error import invalidInputError
  from pyspark.sql.dataframe import DataFrame
  from typing import TYPE_CHECKING, Any, Dict, List, Callable, Union, Optional
label_desc: Use image node_id 13 for calling the SparkTFEstimator obj's underlying member method code with example usage: obj._update_weights(res) and returns: result