pt-resnet-profiling-multi-gpu-single-node.ipynb failed CI

Link to the notebook:
https://github.com/aws/amazon-sagemaker-examples/blob/master/sagemaker-debugger/pytorch_profiling/pt-resnet-profiling-multi-gpu-single-node.ipynb

Error:
Kernel died

Open sagemaker-debugger/pytorch_profiling/pt-resnet-profiling-multi-gpu-single-node.ipynb and reproduce the CI failure that ends with "Kernel died." Investigate the notebook's execution in the CI environment; done means the notebook completes successfully without the kernel dying.

Written by the indexing model from the issue text.