relax exception type

no qa
aws · Aug 16, 2022 · 2d338ca · 2d338ca
1 parent 2104fe4
commit 2d338ca
Showing 1 changed file with 5 additions and 2 deletions.
diff --git a/src/sagemaker_training/smdataparallel.py b/src/sagemaker_training/smdataparallel.py
@@ -33,8 +33,11 @@
 
     # list of exceptions SMDDP wants training toolkit to catch and log
     exception_classes = [x for x in dir(exceptions) if isclass(getattr(exceptions, x))]
-except ImportError:
-    logger.info("No exception classes found in smdistributed.dataparallel")
+# relaxed exception type in case of custom exceptions thrown during import
+except Exception:  # pylint: disable=broad-except
+    logger.info("smdistributed.dataparallel not found or " \
+                "using an older version without custom exceptions." \
+                "SM training toolkit will track user script error only")
     exception_classes = [errors.ExecuteUserScriptError]