Minor fix related to tensorboard in PyTorch

flyteorg · Jul 20, 2021 · e009a19 · e009a19
1 parent ab4092d
commit e009a19
Show file tree

Hide file tree

Showing 2 changed files with 2 additions and 1 deletion.
diff --git a/cookbook/case_studies/ml_training/pytorch/single_node/gpu_training.py b/cookbook/case_studies/ml_training/pytorch/single_node/gpu_training.py
@@ -24,7 +24,7 @@
 from flytekit import Resources, task, workflow
 from flytekit.types.directory import TensorboardLogs
 from flytekit.types.file import PNGImageFile, PythonPickledFile
-from tensorboardX import SummaryWriter
+from torch.utils.tensorboard import SummaryWriter
 from torch import distributed as dist
 from torch import nn, optim
 from torchvision import datasets, transforms

diff --git a/cookbook/case_studies/ml_training/pytorch/single_node/requirements.in b/cookbook/case_studies/ml_training/pytorch/single_node/requirements.in
@@ -1,3 +1,4 @@
 -r ../../../../common/requirements-common.in
+tensorboard
 torch
 torchvision