code: def test_save_and_load_model(only_state_dict):
      """
      测试 save_model 和 load_model 函数
      """
      try:
          path = "model"
          dataset = PaddleNormalXYDataset(20)
          dataloader = DataLoader(dataset, batch_size=4)
          driver1, driver2 = generate_random_driver(
              20, 1, device="gpu"
          ), generate_random_driver(20, 1, device="gpu")
  
          if only_state_dict:
              driver1.save_model(path, only_state_dict)
          else:
              driver1.save_model(
                  path,
                  only_state_dict,
                  input_spec=[paddle.ones((32, 10))],
              )
          driver2.load_model(path, only_state_dict)
  
          for batch in dataloader:
              batch = driver1.move_data_to_device(batch)
              res1 = driver1.model.evaluate_step(**batch)
              res2 = driver2.model.evaluate_step(**batch)
  
              assert paddle.equal_all(res1["pred"], res2["pred"])
      finally:
          if only_state_dict:
              rank_zero_rm(path)
          else:
              rank_zero_rm(path + ".pdiparams")
              rank_zero_rm(path + ".pdiparams.info")
              rank_zero_rm(path + ".pdmodel")
  
imports: import pytest
  from pathlib import Path
  from fastNLP.core.drivers.paddle_driver.single_device import PaddleSingleDriver
  from fastNLP.core.samplers import ReproduceBatchSampler, RandomSampler
  from tests.helpers.models.paddle_model import PaddleNormalModel_Classification_1
  from tests.helpers.datasets.paddle_data import PaddleNormalDataset, PaddleNormalXYDataset
  from tests.helpers.datasets.torch_data import TorchNormalDataset
  from tests.helpers.models.torch_model import TorchNormalModel_Classification_1
  from fastNLP.envs.distributed import rank_zero_rm
  from fastNLP.envs.imports import _NEED_IMPORT_PADDLE, _NEED_IMPORT_TORCH
  from fastNLP import prepare_paddle_dataloader, BucketedBatchSampler
label_desc: Use image node_id 2 for calling a global function with example usage: test_save_and_load_model(only_state_dict) without return types