mirror of
https://github.com/hpcaitech/ColossalAI.git
synced 2025-06-28 08:17:57 +00:00
Fixed import bug for no-tensorboard environment (#354)
This commit is contained in:
parent
6310fb5aae
commit
4ac58ac898
@ -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):
|
||||||
|
Loading…
Reference in New Issue
Block a user