code: def test_gptj_state_dict(model_name):
      config = gptj_config_to_gpt2_config(
          GPTJConfig.from_pretrained(model_name)
      )
      pretrained_state_dict = remap_state_dict_hf_gptj(
          state_dict_from_pretrained(model_name), config
      )
      model = GPTLMHeadModel(
          config, device="meta"
      )  # Without device='meta' init is very slow
      state_dict = model.state_dict()
      assert state_dict.keys() == pretrained_state_dict.keys()
      for k in state_dict.keys():
          assert state_dict[k].shape == pretrained_state_dict[k].shape
  
file_path: flash-attention/tests/models/test_gptj.py
imports: import time
  import pytest
  import torch
  from flash_attn.models.gpt import GPTLMHeadModel
  from flash_attn.models.gptj import gptj_config_to_gpt2_config, remap_state_dict_hf_gptj
  from flash_attn.utils.generation import update_graph_cache
  from flash_attn.utils.pretrained import state_dict_from_pretrained
  from transformers import AutoTokenizer, GPTJConfig
  from transformers.models.gptj.modeling_gptj import GPTJForCausalLM
label_desc: Use image node_id 1 for calling a global function with example usage: test_gptj_state_dict(model_name) without return types