code: def test_get_set_head_states(self):
      config = copy.deepcopy(self._get_config(self.model_configs[0]))
      head_configs = config["model"]["heads"]
      config["model"]["heads"] = []
      model = build_model(config["model"])
      trunk_state = model.get_classy_state()
  
      heads = defaultdict(list)
      for head_config in head_configs:
          head = build_head(head_config)
          heads[head_config["fork_block"]].append(head)
      model.set_heads(heads)
      model_state = model.get_classy_state()
  
      # the heads should be the same as we set
      self.assertEqual(len(heads), len(model.get_heads()))
      for block_name, hs in model.get_heads().items():
          self.assertEqual(hs, heads[block_name])
  
      model.clear_heads()
      self._compare_model_state(model.get_classy_state(), trunk_state)
  
      model.set_heads(heads)
      self._compare_model_state(model.get_classy_state(), model_state)
  
file_path: ClassyVision/test/manual/models_classy_vision_model_test.py
imports: import copy
  import unittest
  from collections import defaultdict
  import torch
  from classy_vision.heads import build_head
  from classy_vision.models import build_model, ClassyModel
  from test.generic.config_utils import get_test_model_configs
  from test.generic.utils import compare_model_state
label_desc: Use image node_id 5 for calling the TestClassyModel obj's underlying member method code with example usage: obj.test_get_set_head_states() without return types