Minor fix related to tensorboard in PyTorch (#334)

Signed-off-by: Jinserk Baik <[email protected]>
flyteorg · Jul 27, 2021 · b9fbdb8 · b9fbdb8
1 parent d0302a0
commit b9fbdb8
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