Fix illegal tb_logger use in distributed training

2020-07-23 09:14:01 -06:00 · 2020-07-23 09:14:01 -06:00 · a7541b6d8d
commit a7541b6d8d
parent bba283776c
1 changed files with 2 additions and 1 deletions
--- a/codes/train.py
+++ b/codes/train.py
@ -203,6 +203,7 @@ def main():
                message += ')] '
                for k, v in logs.items():
                    if 'histogram' in k:
+                        if rank <= 0:
                            tb_logger.add_histogram(k, v, current_step)
                    else:
                        message += '{:s}: {:.4e} '.format(k, v)