kge_pytorch.ipynb failed CI

Link to the notebook:
https://github.com/aws/amazon-sagemaker-examples/blob/master/sagemaker-python-sdk/dgl_kge/kge_pytorch.ipynb

Error:

---------------------------------------------------------------------------
Exception encountered at "In [3]":
---------------------------------------------------------------------------
UnexpectedStatusException Traceback (most recent call last)
in
----> 1 estimator.fit()

/opt/conda/lib/python3.6/site-packages/sagemaker/estimator.py in fit(self, inputs, wait, logs, job_name, experiment_config)
670 self.jobs.append(self.latest_training_job)
671 if wait:
--> 672 self.latest_training_job.wait(logs=logs)
673
674 def _compilation_job_name(self):

/opt/conda/lib/python3.6/site-packages/sagemaker/estimator.py in wait(self, logs)
1601 # If logs are requested, call logs_for_jobs.
1602 if logs != "None":
-> 1603 self.sagemaker_session.logs_for_job(self.job_name, wait=True, log_type=logs)
1604 else:
1605 self.sagemaker_session.wait_for_

[...]

File "/

Start with sagemaker-python-sdk/dgl_kge/kge_pytorch.ipynb and inspect cell In [3], where estimator.fit() fails during CI. Use the shown UnexpectedStatusException as the starting point for reproducing the failure; done means the notebook completes its CI run without this error.

Written by the indexing model from the issue text.