Change tensorboard imports to use built-in tensorboard if available

2019-10-08 08:14:08 -05:00
parent d688af19e5
commit 5ce8d29abe
7 changed files with 41 additions and 10 deletions
--- a/examples/run_multiple_choice.py
+++ b/examples/run_multiple_choice.py
@@ -29,7 +29,12 @@ import torch
 from torch.utils.data import (DataLoader, RandomSampler, SequentialSampler,
                              TensorDataset)
 from torch.utils.data.distributed import DistributedSampler
-from tensorboardX import SummaryWriter
+
+try:
+    from torch.utils.tensorboard import SummaryWriter
+except:
+    from tensorboardX import SummaryWriter
+
 from tqdm import tqdm, trange

 from transformers import (WEIGHTS_NAME, BertConfig,