code: def test_single(self, shuffle, drop_last):
      before_batch_size = 7
      num_batch_per_bucket = 4  # 那么任意 batch 内的长度差值不应该超过4
  
      dataset = DatasetWithVaryLength(num_of_data=1000)
      re_batchsampler = BucketedBatchSampler(
          dataset,
          length=dataset.data,
          batch_size=before_batch_size,
          num_batch_per_bucket=num_batch_per_bucket,
          drop_last=drop_last,
          shuffle=shuffle,
      )
      re_batchsampler.set_epoch(0)
      forward_steps = 10
      iterator = iter(re_batchsampler)
      already_generate_indices = set()
      for _ in range(forward_steps):
          batch = next(iterator)
          assert (
              max(batch) - min(batch)
              <= before_batch_size * num_batch_per_bucket
          )
          already_generate_indices.update(batch)
  
      # 1. 保存状态
      state = re_batchsampler.state_dict()
  
      # 2. 断点重训，继续训练
      re_batchsampler2 = BucketedBatchSampler(
          dataset,
          length=dataset.data,
          batch_size=before_batch_size,
          num_batch_per_bucket=num_batch_per_bucket,
          drop_last=drop_last,
          shuffle=shuffle,
      )
      re_batchsampler2.load_state_dict(state)
      re_batchsampler2.set_epoch(0)
      new_already_generate_indices = set()
      mask = np.ones(len(dataset), dtype=bool)
      mask[list(already_generate_indices)] = 0
      indices = np.arange(len(dataset))[mask]
      max_diff = -1
      for i in range(
          len(indices) - before_batch_size * num_batch_per_bucket
      ):
          max_diff = max(
              max_diff,
              indices[i + before_batch_size * num_batch_per_bucket]
              - indices[i],
          )
      for batch in re_batchsampler2:
          assert max(batch) - min(batch) <= max_diff
          for b in batch:
              assert b not in already_generate_indices
          new_already_generate_indices.update(batch)
      if drop_last is False:
          assert len(
              new_already_generate_indices.union(
                  already_generate_indices
              )
          ) == len(dataset)
  
      # 改变 batch_size；
      after_batch_size = 3
      re_batchsampler3 = BucketedBatchSampler(
          dataset,
          length=dataset.data,
          batch_size=after_batch_size,
          num_batch_per_bucket=num_batch_per_bucket,
          drop_last=drop_last,
          shuffle=shuffle,
      )
      re_batchsampler3.load_state_dict(state)
      re_batchsampler3.set_epoch(0)
      count = 0
  
      mask = np.ones(len(dataset), dtype=bool)
      mask[list(already_generate_indices)] = 0
      indices = np.arange(len(dataset))[mask]
      max_diff = -1
      for i in range(
          len(indices) - after_batch_size * num_batch_per_bucket
      ):
          max_diff = max(
              max_diff,
              indices[i + after_batch_size * num_batch_per_bucket]
              - indices[i],
          )
  
      for batch in re_batchsampler3:
          assert max(batch) - min(batch) <= max_diff
          for b in batch:
              assert b not in already_generate_indices
          already_generate_indices.update(batch)
          count += 1
          if count > 5:
              break
  
      # 再 save ，不允许再上个epoch没结束继续sample
      after_batch_size = 5
      with pytest.raises(RuntimeError):
          state = re_batchsampler3.state_dict()
  
      for batch in re_batchsampler3:  # consume all, 这样才能save
          pass
  
      already_generate_indices = set()
      count = 0
      for batch in re_batchsampler3:  # 重新开始
          assert max(batch) - min(batch) <= max_diff
          for b in batch:
              assert b not in already_generate_indices
          already_generate_indices.update(batch)
          count += 1
          if count > 5:
              break
  
      state = re_batchsampler3.state_dict()
      # 这里的 drop_last 为 False，需要最终是所有 sample
      re_batchsampler4 = BucketedBatchSampler(
          dataset,
          length=dataset.data,
          batch_size=after_batch_size,
          num_batch_per_bucket=num_batch_per_bucket,
          drop_last=False,
          shuffle=shuffle,
      )
      re_batchsampler4.load_state_dict(state)
      re_batchsampler4.set_epoch(0)
  
      mask = np.ones(len(dataset), dtype=bool)
      mask[list(already_generate_indices)] = 0
      indices = np.arange(len(dataset))[mask]
      max_diff = -1
      for i in range(
          len(indices) - after_batch_size * num_batch_per_bucket
      ):
          max_diff = max(
              max_diff,
              indices[i + after_batch_size * num_batch_per_bucket]
              - indices[i],
          )
  
      for batch in re_batchsampler4:
          assert max(batch) - min(batch) <= max_diff
          for b in batch:
              assert b not in already_generate_indices
          already_generate_indices.update(batch)
  
      assert len(already_generate_indices) == len(dataset)
  
imports: import numpy
  import pytest
  from itertools import chain
  from copy import deepcopy
  from array import array
  from tests.helpers.datasets.normal_data import NormalSampler, NormalBatchSampler
  from fastNLP.core.samplers import ReproduceBatchSampler, BucketedBatchSampler, RandomBatchSampler
label_desc: Use image node_id 2 for calling the TestBucketedBatchSampler obj's underlying member method code with example usage: obj.test_single(shuffle, drop_last) without return types