fix: add condition to avoid error when 'model_dir' is None (#402)

aws · Oct 14, 2020 · d2e8848 · d2e8848
1 parent ae8703c
commit d2e8848
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/src/sagemaker_tensorflow_container/training.py b/src/sagemaker_tensorflow_container/training.py
@@ -204,7 +204,7 @@ def _log_model_missing_warning(model_dir):
 
 
 def _model_dir_with_training_job(model_dir, job_name):
-    if model_dir.startswith("/opt/ml"):
+    if model_dir and model_dir.startswith("/opt/ml"):
         return model_dir
     else:
         return "{}/{}/model".format(model_dir, job_name)