code: def _compare_model_state(self, state_1, state_2, check_heads=True):
      return compare_model_state(
          self, state_1, state_2, check_heads=check_heads
      )
  
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 2 for calling the TestUpdateStateFunctions obj's underlying member method code with example usage: obj._compare_model_state(state_1, state_2, check_heads) and returns: compare_model_state