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
index 23870c0562..1058cdc8c8 100644
--- 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
index 5b84192d70..d88480c387 100644
--- 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