mirror of https://github.com/InternLM/InternLM
Update train.py
parent
f8eaf618af
commit
f5e65c0f30
3
train.py
3
train.py
|
@ -192,10 +192,9 @@ def main(args):
|
||||||
# transfer the train data loader into train data iterator
|
# transfer the train data loader into train data iterator
|
||||||
train_iter = iter(train_dl)
|
train_iter = iter(train_dl)
|
||||||
|
|
||||||
|
with initialize_llm_profile(profiling=args.profiling, start_time=current_time) as prof:
|
||||||
# close automatic garbage collection
|
# close automatic garbage collection
|
||||||
gc.disable()
|
gc.disable()
|
||||||
|
|
||||||
with initialize_llm_profile(profiling=args.profiling, start_time=current_time) as prof:
|
|
||||||
# start iterating the train data and begin training
|
# start iterating the train data and begin training
|
||||||
for batch_count in range(train_state.batch_count, total_steps):
|
for batch_count in range(train_state.batch_count, total_steps):
|
||||||
empty_cache_and_diag(batch_count, interval=gpc.config.data.empty_cache_and_diag_interval)
|
empty_cache_and_diag(batch_count, interval=gpc.config.data.empty_cache_and_diag_interval)
|
||||||
|
|
Loading…
Reference in New Issue