code: def test_save_and_load_with_randomsampler(
      only_state_dict, use_dataloader
  ):
      """
      测试save和load函数，主要测试 dataloader 被替换了 sampler 的情况
      """
  
      try:
          path = "model.ckp"
  
          driver1, driver2 = generate_random_driver(
              20, 1, device="gpu", lr=0.01
          ), generate_random_driver(20, 1, device="gpu", lr=0.001)
          dataset = JittorNormalXYDataset(20)
          dataloader = get_dataloader(
              dataset,
              use_dataloader,
              sampler=RandomSampler(dataset, True),
              batch_size=4,
              shuffle=True,
          )
          num_consumed_batches = 2
  
          already_seen_x_set = set()
          already_seen_y_set = set()
          driver1.set_sampler_epoch(dataloader, 7)
          for idx, batch in enumerate(dataloader):
              if idx >= num_consumed_batches:
                  break
              already_seen_x_set.update(
                  batch["x"]
                  .reshape(
                      -1,
                  )
                  .tolist()
              )
              already_seen_y_set.update(
                  batch["y"]
                  .reshape(
                      -1,
                  )
                  .tolist()
              )
  
          sampler_states = dataloader.sampler.state_dict()
          save_states = {"num_consumed_batches": num_consumed_batches}
          driver1.save_checkpoint(
              Path(path),
              save_states,
              dataloader,
              only_state_dict,
              should_save_model=True,
          )
  
          # 加载
          # 更改 batch_size
          dataloader = get_dataloader(
              dataset,
              use_dataloader,
              sampler=RandomSampler(dataset, True),
              batch_size=2,
              shuffle=True,
          )
          load_states = driver2.load_checkpoint(
              Path(path),
              dataloader,
              only_state_dict,
              should_load_model=True,
          )
          replaced_loader = load_states.pop("dataloader")
  
          # 1. 检查 optimizer 的状态
          assert driver2.optimizers[0].lr == driver1.optimizers[0].lr
  
          # 2. 检查 sampler 是否被正确地加载和替换
          assert not (replaced_loader is dataloader)
          assert isinstance(replaced_loader.sampler, RandomSampler)
          assert replaced_loader.sampler.seed == sampler_states["seed"]
          assert (
              replaced_loader.sampler.epoch == sampler_states["epoch"]
          )
          assert (
              replaced_loader.sampler.num_consumed_samples
              == 4 * num_consumed_batches
          )
          assert (
              replaced_loader.sampler.dataset.total_len
              == sampler_states["length"]
          )
          assert (
              replaced_loader.sampler.shuffle
              == sampler_states["shuffle"]
          )
  
          # 4. 检查 model 的参数是否正确
          # 5. 检查 batch_idx
          start_batch = load_states.pop("batch_idx_in_epoch")
          assert start_batch == 2 * num_consumed_batches
          left_x_batches = set()
          left_y_batches = set()
          driver2.set_sampler_epoch(replaced_loader, 7)
          for idx, batch in enumerate(replaced_loader):
              left_x_batches.update(
                  batch["x"]
                  .reshape(
                      -1,
                  )
                  .tolist()
              )
              left_y_batches.update(
                  batch["y"]
                  .reshape(
                      -1,
                  )
                  .tolist()
              )
              res1 = driver1.model.evaluate_step(**batch)
              res2 = driver2.model.evaluate_step(**batch)
              assert jt.all_(jt.equal(res1["pred"], res2["pred"]))
  
          assert (
              len(left_x_batches) + len(already_seen_x_set)
              == dataset.total_len
          )
          assert (
              len(left_x_batches | already_seen_x_set)
              == dataset.total_len
          )
          assert (
              len(left_y_batches) + len(already_seen_y_set)
              == dataset.total_len
          )
          assert (
              len(left_y_batches | already_seen_y_set)
              == dataset.total_len
          )
      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 4 for calling a global function with example usage: test_save_and_load_with_randomsampler(only_state_dict, use_dataloader) without return types