From 534e0bb11801b40c7cb40298ac374b5ae24f80ee Mon Sep 17 00:00:00 2001
From: HELSON <72907851+1SAA@users.noreply.github.com>
Date: Wed, 9 Mar 2022 19:48:04 +0800
Subject: [PATCH] Fixed import bug for no-tensorboard environment (#354)

---
 colossalai/utils/profiler/comm_profiler.py | 3 +--
 1 file changed, 1 insertion(+), 2 deletions(-)

diff --git a/colossalai/utils/profiler/comm_profiler.py b/colossalai/utils/profiler/comm_profiler.py
index 25538dff0..b2356e54e 100644
--- a/colossalai/utils/profiler/comm_profiler.py
+++ b/colossalai/utils/profiler/comm_profiler.py
@@ -5,7 +5,6 @@ import torch
 from torch.autograd.profiler import profile
 import torch.distributed as dist
 from torch.distributed import ReduceOp
-import torch.utils.tensorboard as tb
 from colossalai.utils import get_current_device
 from .prof_utils import BaseProfiler
 from typing import List, Optional
@@ -126,7 +125,7 @@ class CommProfiler(BaseProfiler):
         dist.broadcast = torch_broadcast
         dist.reduce = torch_reduce
 
-    def to_tensorboard(self, writer: tb.writer):
+    def to_tensorboard(self, writer):
         writer.add_text(tag="Collective Communication", text_string=self.result_list("\n\n"))
 
     def to_file(self, filename: Path):