code: def check_models(
      original_unwrapped_model: torch.nn.Module,
      model_after_unwrapping: torch.nn.Module,
  ) -> None:
      """Helper function to test original and unwrapped models are the same."""
      for layer_id in range(len(original_unwrapped_model)):
          # for unwrapped models, we should be able to index them
          assert repr(model_after_unwrapping[layer_id]) == repr(
              original_unwrapped_model[layer_id]
          )
  
imports: import os
  import torch
  import torch.distributed
  from utils.common_utils import unwrap_model_fn
label_desc: Use image node_id 1 for calling a global function with example usage: check_models(original_unwrapped_model, model_after_unwrapping) without return types