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

This commit is contained in:
HELSON 2022-03-09 19:48:04 +08:00 committed by GitHub
parent 6310fb5aae
commit 4ac58ac898
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23

View File

@ -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):
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):