Update train.py

This commit is contained in:
Mingshuang Luo 2021-10-12 20:16:51 +08:00 committed by GitHub
parent 18d7dd2696
commit 906d0ad6f8
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23

View File

@ -435,7 +435,7 @@ def train_one_epoch(
f"batch {batch_idx}, loss[{loss_info}], "
f"tot_loss[{tot_loss}], batch size: {batch_size}"
)
if batch_idx % 10 == 0:
if batch_idx % params.log_interval == 0:
if tb_writer is not None:
loss_info.write_summary(