[🐛BUG] GradientState' object has no attribute '_set_end_of_dataloader

[minji-o-j]

Describe the bug GradientState' object has no attribute '_set_end_of_dataloader

To Reproduce python run_textbox.py --model=PTG --dataset=xsum --model_path=facebook/bart-large

Logs If applicable, add the project.log of your project in the saved folder to help us make a better diagnosis.

train    1:   0%|                                                                                                                   | 1/1063 [00:22<6:35:55, 22.37s/step, loss=4.86]06 Jul 06:43    ERROR Traceback (most recent call last):
  File "/workspace/TextBox/textbox/utils/dashboard.py", line 310, in new_experiment
    yield True
  File "/workspace/TextBox/textbox/quick_start/experiment.py", line 127, in run
    self._do_train_and_valid()
  File "/workspace/TextBox/textbox/quick_start/experiment.py", line 104, in _do_train_and_valid
    self.valid_result = self.trainer.fit(train_data, valid_data)
  File "/workspace/TextBox/textbox/trainer/trainer.py", line 423, in fit
    loss = self._train_epoch(train_data, epoch_idx, valid_data)["loss"]
  File "/workspace/TextBox/textbox/trainer/trainer.py", line 223, in _train_epoch
    self.accelerator.gradient_state._set_end_of_dataloader(False)
AttributeError: 'GradientState' object has no attribute '_set_end_of_dataloader'

[minji-o-j]

https://github.com/ailab-prompt-transfer/TextBox/blob/0d7debd9817c0f0278a52506e41ac2677d567550/textbox/trainer/trainer.py#L236C30-L236C30

  if self.accelerator.sync_gradients:
      if not self.disable_tqdm:
          train_tqdm.update(1)
          train_tqdm.set_postfix(loss=self._summary_tracker.epoch_loss)
      if valid_data:
          self.stopped |= self._valid(valid_data, "step")
          # self.accelerator.gradient_state._set_end_of_dataloader(False) # << delete
      if self.stopped:
          break