Num_of_step increase infinitely but num_of_epoch stuck 0 with dynamic batch size dataset

[Lingeng56]

Hey, I implemented a dynamic batch size dataset implement data batching procedure in dataset class rather than using sampler.
My dataset
(Functions <shuffle, sort, batch> in __batched_data will all return a generator)

class CustomDataset(IterableDataset):

    def __init__(self, data_config, mode):
        self.data_config = data_config
        self.mode = mode
        self.__prepare_data()
        self.__batched_data()


    def __batched_data(self):
        if self.data_config['shuffle']:
            self.dataset = shuffle(self.dataset,
                                   self.data_config['shuffle_size'])

        if self.data_config['sort']:
            self.dataset = sort(self.dataset,
                                self.data_config['sort_size'])

        self.dataset = batch(self.dataset,
                             self.data_config['batch_type'],
                             batch_size=self.data_config['batch_size'],
                             max_frames_in_batch=self.data_config['max_frames_in_batch'])

        self.dataset = padding(self.dataset)

    def __prepare_data(self):
        if os.path.exists(self.data_config[f'{self.mode}_pkl_path']):
            self.dataset = pickle.load(open(self.data_config[f'{self.mode}_pkl_path'], 'rb'))
        else:
            dataset = DataList(self.data_config[f'{self.mode}_data_list_path'])
            self.dataset = []
            for sample in tqdm(dataset, desc='Preprocessing Data...'):
                sample = parse_raw(sample)
                sample = tokenize(sample,
                                  vocabs=self.data_config['vocabs'],
                                  bpe_model=self.data_config['bpe_model'],
                                  non_lang_syms=self.data_config['non_lang_syms'],
                                  split_with_space=self.data_config['split_with_space']
                                  )
                check = filter_data(sample,
                                    max_length=self.data_config['max_length'],
                                    min_length=self.data_config['min_length'],
                                    token_max_length=self.data_config['token_max_length'],
                                    token_min_length=self.data_config['token_min_length'],
                                    min_output_input_ratio=self.data_config['min_output_input_ratio'],
                                    max_output_input_ratio=self.data_config['max_output_input_ratio'])

                if not check:
                    continue

                sample = resample(sample,
                                  resample_rate=self.data_config['resample_rate'])

                if self.data_config['speed_perturb']:
                    sample = speed_perturb(sample,
                                           speeds=self.data_config['speeds'])

                if self.data_config['feat_type'] == 'fbank':
                    sample = compute_fbank(sample,
                                           num_mel_bins=self.data_config['num_mel_bins'],
                                           frame_length=self.data_config['frame_length'],
                                           frame_shift=self.data_config['frame_shift'],
                                           dither=self.data_config['dither'])
                elif self.data_config['feat_type'] == 'mfcc':
                    sample = compute_mfcc(sample,
                                          num_mel_bins=self.data_config['num_mel_bins'],
                                          frame_length=self.data_config['frame_length'],
                                          dither=self.data_config['dither'],
                                          num_ceps=self.data_config['num_ceps'],
                                          high_freq=self.data_config['high_freq'],
                                          low_freq=self.data_config['low_freq'])

                if self.data_config['spec_aug']:
                    sample = spec_aug(sample,
                                      num_t_mask=self.data_config['num_t_mask'],
                                      num_f_mask=self.data_config['num_f_mask'],
                                      max_t=self.data_config['max_t'],
                                      max_f=self.data_config['max_f'],
                                      )
                self.dataset.append(sample)
            pickle.dump(self.dataset, open(self.data_config[f'{self.mode}_pkl_path'], 'wb'))
        print(f'Loaded {self.mode} Dataset ({len(self.dataset)} Samples)')

    def __iter__(self):
        for sample in self.dataset:
            yield sample

    def __len__(self):
        return len(self.dataset)

My dataloader

train_dataloader = DataLoader(train_dataset,
                                  batch_size=None,
                                  num_workers=args.num_workers,
                                  pin_memory=args.pin_memory,
                                  prefetch_factor=args.prefetch)

During training, I find my progress bar can't get the total number of steps of an epoch, and this was to my expected. But during training, number of step increased infinitely and epoch stuck to 0.

Maybe anyone faced same problem, and could give me some suggestions...