class_docstr: 测试多卡情况下 save 和 load 相关函数的表现
class_docstr_tok: ["\u6d4b\u8bd5\u591a\u5361\u60c5\u51b5\u4e0b","save","\u548c","load","\u76f8\u5173\u51fd\u6570\u7684\u8868\u73b0"]
docstr: 测试 save_model 和 load_model 函数
docstr_tok: ["\u6d4b\u8bd5","save_model","\u548c","load_model","\u51fd\u6570"]
code: def test_save_and_load_model(self, only_state_dict):
      """
      测试 save_model 和 load_model 函数
      """
      try:
          path = "model"
  
          dataloader = DataLoader(self.dataset, batch_size=2)
          driver1, driver2 = generate_driver(20, 1), generate_driver(
              20, 1
          )
  
          driver1.save_model(path, only_state_dict)
  
          # 同步
          comm.barrier()
          driver2.load_model(path, only_state_dict)
  
          for idx, batch in enumerate(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 os
  import sys
  import pytest
  from pathlib import Path
  from fastNLP.core.drivers.oneflow_driver.ddp import OneflowDDPDriver
  from fastNLP import prepare_oneflow_dataloader
  from fastNLP.core.samplers import RandomSampler, UnrepeatedSampler, BucketedBatchSampler, UnrepeatedRandomSampler, UnrepeatedSequentialSampler
  from tests.helpers.models.oneflow_model import OneflowNormalModel_Classification_1
  from tests.helpers.datasets.oneflow_data import OneflowNormalDataset, OneflowNormalXYDataset
  from tests.helpers.utils import recover_logger
  from fastNLP.envs.distributed import rank_zero_rm
  from fastNLP import logger
  from fastNLP.core.drivers.oneflow_driver.dist_utils import fastnlp_oneflow_all_gather
  from fastNLP.envs.imports import _NEED_IMPORT_ONEFLOW
label_desc: Use image node_id 2 for calling the TestSaveLoad obj's underlying member method code with example usage: obj.test_save_and_load_model(only_state_dict) without return types