code: def _assert_checkpoint_equal(model, output_path):
      assert os.path.exists(output_path)
      single_output = torch.load(output_path)
      state_dict = model.state_dict()
      for orig_param, saved_model_param in zip(
          state_dict.values(), single_output["state_dict"].values()
      ):
          if model.dtype == torch.half:
              # moved model to float32 for comparison with single fp32 saved weights
              saved_model_param = saved_model_param.half()
          assert torch.equal(orig_param.cpu(), saved_model_param)
  
file_path: pytorch-lightning/tests/tests_pytorch/utilities/test_deepspeed_collate_checkpoint.py
imports: import os
  import torch
  from lightning.pytorch import Trainer
  from lightning.pytorch.demos.boring_classes import BoringModel
  from lightning.pytorch.strategies import DeepSpeedStrategy
  from lightning.pytorch.utilities.deepspeed import convert_zero_checkpoint_to_fp32_state_dict
  from tests_pytorch.helpers.runif import RunIf
label_desc: Use image node_id 2 for calling a global function with example usage: _assert_checkpoint_equal(model, output_path) without return types