Skip to content

Commit

Permalink
fix: import error when missing tensorflow (now optional)
Browse files Browse the repository at this point in the history
  • Loading branch information
sehoffmann committed Jan 5, 2025
1 parent aad274d commit 97b3642
Showing 1 changed file with 6 additions and 1 deletion.
7 changes: 6 additions & 1 deletion dmlcloud/core/callbacks.py
Original file line number Diff line number Diff line change
Expand Up @@ -9,7 +9,6 @@
import torch
from omegaconf import OmegaConf
from progress_table import ProgressTable
from torch.utils.tensorboard import SummaryWriter

from ..git import git_diff
from ..util.logging import DevNullIO, experiment_header, general_diagnostics, IORedirector
Expand Down Expand Up @@ -426,8 +425,14 @@ class TensorboardCallback(Callback):

def __init__(self, log_dir: Union[str, Path]):
self.log_dir = Path(log_dir)
try:
from torch.utils.tensorboard import SummaryWriter # noqa: F401
except ImportError:
raise ImportError('tensorflow is required for the TensorboardCallback')

def pre_run(self, pipe):
from torch.utils.tensorboard import SummaryWriter

self.writer = SummaryWriter(log_dir=self.log_dir)

def post_epoch(self, stage: 'Stage'):
Expand Down

0 comments on commit 97b3642

Please sign in to comment.