diff --git a/assets/model_monitoring/components/src/model_monitor_feature_selector/run.py b/assets/model_monitoring/components/src/model_monitor_feature_selector/run.py
index e63fd03f45..241491ef9a 100644
--- a/assets/model_monitoring/components/src/model_monitor_feature_selector/run.py
+++ b/assets/model_monitoring/components/src/model_monitor_feature_selector/run.py
@@ -13,7 +13,7 @@
 from shared_utilities.io_utils import (
     try_read_mltable_in_spark, try_read_mltable_in_spark_with_error, save_spark_df_as_mltable
 )
-from shared_utilities.momo_exceptions import DataNotFoundError, InvalidInputError
+from shared_utilities.momo_exceptions import InvalidInputError
 
 
 def run():
@@ -42,16 +42,12 @@ def run():
         input_df1 = input_df2
 
     feature_importance = None
-    try:
-        feature_importance = try_read_mltable_in_spark_with_error(args.feature_importance, "feature_importance")
-    except DataNotFoundError as e:
-        if args.filter_type == FeatureSelectorType.TOP_N_BY_ATTRIBUTION.name:
-            raise e
-    except Exception:
-        if args.filter_type == FeatureSelectorType.TOP_N_BY_ATTRIBUTION.name:
-            raise Exception(
-                "Error encountered when retrieving top N features. Please ensure target_column is defined."
-            )
+    if args.filter_type == FeatureSelectorType.TOP_N_BY_ATTRIBUTION.name:
+        try:
+            feature_importance = try_read_mltable_in_spark_with_error(args.feature_importance, "feature_importance")
+        except Exception as e:
+            raise Exception(f"Error encountered when retrieving top N features. {e}. "
+                            "The feature importance step may not have completed successfully.")
 
     feature_selector = FeatureSelectorFactory().produce(
         feature_selector_type=args.filter_type,