docstr: Tests that the update_classy_model successfully updates from a
  checkpoint
docstr_tok: ["Tests","that","the","update_classy_model","successfully","updates","from","a","checkpoint"]
code: def test_update_classy_model(self):
      """
      Tests that the update_classy_model successfully updates from a
      checkpoint
      """
      config = get_fast_test_task_config()
      task = build_task(config)
      trainer = LocalTrainer()
      trainer.train(task)
      for reset_heads in [False, True]:
          task_2 = build_task(config)
          # prepare task_2 for the right device
          task_2.prepare()
          update_classy_model(
              task_2.model,
              task.model.get_classy_state(deep_copy=True),
              reset_heads,
          )
          self._compare_model_state(
              task.model.get_classy_state(),
              task_2.model.get_classy_state(),
              check_heads=not reset_heads,
          )
          if reset_heads:
              # the model head states should be different
              with self.assertRaises(Exception):
                  self._compare_model_state(
                      task.model.get_classy_state(),
                      task_2.model.get_classy_state(),
                      check_heads=True,
                  )
  
imports: import copy
  import shutil
  import tempfile
  import unittest
  import unittest.mock
  from pathlib import Path
  import classy_vision.generic.util
  import torch
  import torch.nn
  from classy_vision.generic.util import CHECKPOINT_FILE, get_torch_version, load_checkpoint, save_checkpoint, split_batchnorm_params, Timer, update_classy_model, update_classy_state
  from classy_vision.models import build_model
  from classy_vision.tasks import build_task
  from classy_vision.trainer import LocalTrainer
  from test.generic.config_utils import get_fast_test_task_config, get_test_model_configs
  from test.generic.utils import compare_model_state, compare_states
label_desc: Use image node_id 4 for calling the TestUpdateStateFunctions obj's underlying member method code with example usage: obj.test_update_classy_model() without return types