from local

This commit is contained in:
dohe0342 2023-05-24 21:08:50 +09:00
parent 6fbbacf94c
commit f3cbdad979
2 changed files with 2 additions and 0 deletions

View File

@ -1129,6 +1129,7 @@ def train_one_epoch(
and params.batch_idx_train % params.save_every_n == 0 and params.batch_idx_train % params.save_every_n == 0
): ):
params.cur_batch_idx = batch_idx params.cur_batch_idx = batch_idx
'''
save_checkpoint_with_global_batch_idx( save_checkpoint_with_global_batch_idx(
out_dir=params.exp_dir, out_dir=params.exp_dir,
global_batch_idx=params.batch_idx_train, global_batch_idx=params.batch_idx_train,
@ -1141,6 +1142,7 @@ def train_one_epoch(
scaler=scaler, scaler=scaler,
rank=rank, rank=rank,
) )
'''
del params.cur_batch_idx del params.cur_batch_idx
if rank == 0: if rank == 0: