Skip to content
Open
Changes from 1 commit
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
7 changes: 5 additions & 2 deletions swift/trainers/callback.py
Original file line number Diff line number Diff line change
Expand Up @@ -44,9 +44,12 @@ def add_train_message(logs, state, start_time) -> None:
class ProgressCallbackNew(ProgressCallback):

def on_train_begin(self, args, state, control, **kwargs):
initial_step = state.global_step if state.global_step is not None else 0
if state.is_world_process_zero:
self.training_bar = tqdm(desc='Train', total=state.max_steps, dynamic_ncols=True)
self.current_step = 0
bar_initial = min(initial_step, state.max_steps) if state.max_steps else initial_step
self.training_bar = tqdm(
desc='Train', total=state.max_steps, initial=bar_initial, dynamic_ncols=True)
self.current_step = initial_step
self.start_time = time.time()

def on_prediction_step(self, args, state: TrainerState, control, eval_dataloader=None, **kwargs):
Expand Down
Loading