write the loss information to a tensorboard file. Only for pytorch DDP mode.
(writer, *args)
| 176 | return writer |
| 177 | |
| 178 | def write_tensorboard(writer, *args): |
| 179 | ''' |
| 180 | write the loss information to a tensorboard file. |
| 181 | Only for pytorch DDP mode. |
| 182 | ''' |
| 183 | if dist.get_rank() == 0: # real tensorboard |
| 184 | writer.add_scalar(args[0], args[1], args[2]) |
| 185 | |
| 186 | ################################################################################# |
| 187 | # EMA Update/ DDP Training Utils # |