def METHOD_NAME(self, update_steps: int, train_epochs: int) -> None:
    data = []
    # warmup
    pbar = tqdm(range(update_steps), desc=f'Train epoch [1/{train_epochs}]', disable=not is_rank_0())
    self._on_epoch_start(0)
    self._learn_epoch(pbar, data)
    self._on_epoch_end(0)
    # item is already a batch
    dataloader = DataLoader(data,
                            batch_size=1,
                            shuffle=True,
                            pin_memory=self.dataloader_pin_memory,
                            collate_fn=lambda x: x[0])
    for epoch in range(1, train_epochs):
        pbar = tqdm(dataloader, desc=f'Train epoch [{epoch + 1}/{train_epochs}]', disable=not is_rank_0())
        self._on_epoch_start(epoch)
        self._learn_epoch(pbar, data)
        self._on_epoch_end(epoch)