[Fix]Make sure the args tb_writer passed to the TensorBoardCallback works

[iamlockelightning]

What does this PR do?

Add an if check in the on_train_begin function of the TensorBoardCallback, to make sure the args tb_writer passed to the __init__ won't be overwritten.

I wanted to pass a tb_writer to the TensorBoardCallback, so that I could continue using it after training. But the passed tb_writer didn't work.

tb_writer = SummaryWriter(log_dir="some_dir")
cb = TensorBoardCallback(tb_writer)
trainer = Trainer(
    # not important
    callbacks=[cb],
)
trainer.train()
tb_writer.add_scalar("other metrics", 0.99, 0) # didn't work

After checking the code of TensorBoardCallback, I found that the tb_writer passed to the __init__ would be overwritten by the on_train_begin function with a new one no matter what.

By adding an if statement before the init, this bug should be fixed.

if self.tb_writer is None: # check
    self._init_summary_writer(args, log_dir)

Before submitting

• This PR fixes a typo or improves the docs (you can dismiss the other checks if that's the case).

• Did you read the contributor guideline,
  Pull Request section?

• Was this discussed/approved via a Github issue or the forum? Please add a link
  to it if that's the case.
  Add option to keep tb_writer open after training is done. #5329 (a similar issue but is "wontfix" and out-of-date)

• Did you make sure to update the documentation with your changes? Here are the
  documentation guidelines, and
  here are tips on formatting docstrings.

• Did you write any new necessary tests?

Who can review?

@sgugger

[sgugger]

LGTM, thanks for fixing this!