Commit 534e0bb1 authored by HELSON's avatar HELSON Committed by Frank Lee
Browse files

Fixed import bug for no-tensorboard environment (#354)

parent c57e0898
...@@ -5,7 +5,6 @@ import torch ...@@ -5,7 +5,6 @@ import torch
from torch.autograd.profiler import profile from torch.autograd.profiler import profile
import torch.distributed as dist import torch.distributed as dist
from torch.distributed import ReduceOp from torch.distributed import ReduceOp
import torch.utils.tensorboard as tb
from colossalai.utils import get_current_device from colossalai.utils import get_current_device
from .prof_utils import BaseProfiler from .prof_utils import BaseProfiler
from typing import List, Optional from typing import List, Optional
...@@ -126,7 +125,7 @@ class CommProfiler(BaseProfiler): ...@@ -126,7 +125,7 @@ class CommProfiler(BaseProfiler):
dist.broadcast = torch_broadcast dist.broadcast = torch_broadcast
dist.reduce = torch_reduce 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")) writer.add_text(tag="Collective Communication", text_string=self.result_list("\n\n"))
def to_file(self, filename: Path): def to_file(self, filename: Path):
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment