code: def test_save_and_load_model(only_state_dict):
      """
      测试 save_model 和 load_model 函数
      """
      try:
          path = "model"
          dataset = OneflowNormalXYDataset(20)
          dataloader = DataLoader(dataset, batch_size=4)
          driver1, driver2 = generate_random_driver(
              20, 1
          ), generate_random_driver(20, 1)
  
          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 oneflow.all(res1["pred"] == res2["pred"])
      finally:
          rank_zero_rm(path)
  
imports: import pytest
  from pathlib import Path
  from fastNLP.core.drivers.oneflow_driver.single_device import OneflowSingleDriver
  from fastNLP.core.samplers import ReproduceBatchSampler, RandomSampler
  from tests.helpers.models.oneflow_model import OneflowNormalModel_Classification_1
  from tests.helpers.datasets.oneflow_data import OneflowNormalDataset, OneflowNormalXYDataset
  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_ONEFLOW, _NEED_IMPORT_TORCH
  from fastNLP import prepare_oneflow_dataloader, BucketedBatchSampler
label_desc: Use image node_id 4 for calling a global function with example usage: test_save_and_load_model(only_state_dict) without return types