code: def test_save_and_load_with_randomsampler(only_state_dict, fp16):
      """
      测试save和load函数，主要测试 dataloader 被替换了 sampler 的情况
      """
  
      try:
          path = "model.ckp"
  
          driver1, driver2 = generate_random_driver(
              40, 1, fp16, "gpu"
          ), generate_random_driver(40, 1, False, "gpu")
          dataset = PaddleNormalXYDataset(40)
          batch_sampler = BatchSampler(dataset=dataset, batch_size=4)
          batch_sampler.sampler = RandomSampler(dataset, True)
          dataloader = DataLoader(dataset, batch_sampler=batch_sampler)
          num_consumed_batches = 2
  
          already_seen_x_set = set()
          already_seen_y_set = set()
          driver1.set_sampler_epoch(dataloader, 3)
          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.batch_sampler.sampler.state_dict()
          save_states = {"num_consumed_batches": num_consumed_batches}
          if only_state_dict:
              driver1.save_checkpoint(
                  Path(path),
                  save_states,
                  dataloader,
                  only_state_dict,
                  should_save_model=True,
              )
          else:
              driver1.save_checkpoint(
                  Path(path),
                  save_states,
                  dataloader,
                  only_state_dict,
                  should_save_model=True,
                  input_spec=[paddle.ones((16, 10))],
              )
  
          # 加载
          # 更改 batch_size
          batch_sampler = BatchSampler(dataset=dataset, batch_size=2)
          batch_sampler.sampler = RandomSampler(dataset, True)
          dataloader = DataLoader(dataset, batch_sampler=batch_sampler)
          load_states = driver2.load_checkpoint(
              Path(path),
              dataloader,
              only_state_dict,
              should_load_model=True,
          )
          replaced_loader = load_states.pop("dataloader")
  
          # 1. 检查 optimizer 的状态
          # TODO optimizer 的 state_dict 总是为空
  
          # 2. 检查 sampler 是否被正确地加载和替换
          assert not (replaced_loader is dataloader)
          assert isinstance(
              replaced_loader.batch_sampler.sampler, RandomSampler
          )
          assert (
              replaced_loader.batch_sampler.sampler.seed
              == sampler_states["seed"]
          )
          assert (
              replaced_loader.batch_sampler.sampler.epoch
              == sampler_states["epoch"]
          )
          assert (
              replaced_loader.batch_sampler.sampler.num_consumed_samples
              == 4 * num_consumed_batches
          )
          assert (
              len(replaced_loader.batch_sampler.sampler.dataset)
              == sampler_states["length"]
          )
          assert (
              replaced_loader.batch_sampler.sampler.shuffle
              == sampler_states["shuffle"]
          )
  
          # 3. 检查 fp16 是否被加载
          if fp16:
              assert not isinstance(
                  driver2.grad_scaler, paddle.amp.GradScaler
              )
  
          # 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()
          driver1.set_sampler_epoch(replaced_loader, 3)
          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 paddle.equal_all(res1["pred"], res2["pred"])
  
          assert len(left_x_batches) + len(already_seen_x_set) == len(
              dataset
          )
          assert len(left_x_batches | already_seen_x_set) == len(
              dataset
          )
          assert len(left_y_batches) + len(already_seen_y_set) == len(
              dataset
          )
          assert len(left_y_batches | already_seen_y_set) == len(
              dataset
          )
      finally:
          rank_zero_rm(path)
  
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 4 for calling a global function with example usage: test_save_and_load_with_randomsampler(only_state_dict, fp16) without return types