github/ColossalAI
1
0
Fork 0
mirror of https://github.com/hpcaitech/ColossalAI.git synced 2025-07-18 01:12:42 +00:00
Code Issues Packages Projects Releases Wiki Activity

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

Browse Source
This commit is contained in:
HELSON 2022-03-09 19:48:04 +08:00 committed by Frank Lee
parent c57e089824
commit 534e0bb118
1 changed files with 1 additions and 2 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

3
colossalai/utils/profiler/comm_profiler.py
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):