How to solve the RuntimeError when using torch.utils.tensorboard to add a graph