code: def test_save_and_load_model(only_state_dict, use_dataloader):
      """
      测试 save_model 和 load_model 函数
      """
      try:
          path = "model"
          dataset = JittorNormalXYDataset(20)
          dataloader = get_dataloader(
              dataset,
              sampler=None,
              use_dataloader=use_dataloader,
              batch_size=4,
              shuffle=True,
          )
          driver1, driver2 = generate_random_driver(
              20, 1, device="gpu"
          ), generate_random_driver(20, 1, device="gpu")
  
          driver1.save_model(path, only_state_dict)
          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 jt.all_(jt.equal(res1["pred"], res2["pred"]))
      finally:
          rank_zero_rm(path)
  
imports: import pytest
  from copy import deepcopy
  from pathlib import Path
  from fastNLP.core.drivers.jittor_driver import JittorSingleDriver
  from fastNLP.core.samplers import ReproduceBatchSampler, RandomSampler
  from fastNLP.core.dataloaders import JittorDataLoader
  from tests.helpers.models.jittor_model import JittorNormalModel_Classification_1
  from tests.helpers.datasets.jittor_data import JittorNormalDataset, JittorNormalXYDataset
  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_JITTOR, _NEED_IMPORT_TORCH
label_desc: Use image node_id 3 for calling a global function with example usage: test_save_and_load_model(only_state_dict, use_dataloader) without return types