code: def _check_params(model, expected_param_shapes):
      current_param_shapes = {
          name: tuple(param.shape)
          for name, param in model.named_parameters()
      }
      assert set(current_param_shapes.keys()) == set(
          expected_param_shapes.keys()
      )
      for key, current_shape in current_param_shapes.items():
          expected_shape = expected_param_shapes[key]
          assert (
              current_shape == expected_shape
          ), f"Parameter {key} should have shape {expected_shape}, but found shape {current_shape}"
  
file_path: fairscale/tests/nn/data_parallel/test_fsdp_with_checkpoint_wrapper.py
imports: import contextlib
  import pytest
  import torch
  from torch import nn
  import torch.distributed
  import torch.multiprocessing
  from fairscale.fair_dev.testing.testing import dist_init, skip_if_single_gpu, teardown, temp_files_ctx
  from fairscale.nn.checkpoint.checkpoint_activations import checkpoint_wrapper
  from fairscale.nn.data_parallel import FullyShardedDataParallel
label_desc: Use image node_id 5 for calling a global function with example usage: _check_params(model, expected_param_shapes) without return types