code: def compare_model_state(
      test_fixture, state, state2, check_heads=True
  ):
      for k in state["model"]["trunk"].keys():
          if not torch.allclose(
              state["model"]["trunk"][k], state2["model"]["trunk"][k]
          ):
              print(
                  k,
                  state["model"]["trunk"][k],
                  state2["model"]["trunk"][k],
              )
          test_fixture.assertTrue(
              torch.allclose(
                  state["model"]["trunk"][k],
                  state2["model"]["trunk"][k],
              )
          )
      if check_heads:
          for block, head_states in state["model"]["heads"].items():
              for head_id, states in head_states.items():
                  for k in states.keys():
                      test_fixture.assertTrue(
                          torch.allclose(
                              state["model"]["heads"][block][head_id][
                                  k
                              ],
                              state2["model"]["heads"][block][head_id][
                                  k
                              ],
                          )
                      )
  
imports: import os
  import unittest
  from functools import wraps
  import torch
  from classy_vision.hooks import ClassyHook
  from classy_vision.tasks import ClassyTask
  from classy_vision.trainer import ClassyTrainer
label_desc: Use image node_id 8 for calling a global function with example usage: compare_model_state(test_fixture, state, state2, check_heads) without return types