Minor fix

pomonam · Mar 15, 2024 · 70b70a6 · 70b70a6
1 parent 3b997fb
commit 70b70a6
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/kronfluence/computer/computer.py b/kronfluence/computer/computer.py
@@ -82,14 +82,14 @@ def __init__(
 
         try:
             tracked_module_names = get_tracked_module_names(self.model)
-        except TrackedModuleNotFoundError:
+        except TrackedModuleNotFoundError as e:
             error_msg = (
                 f"No tracked modules found in the provided model: {self.model}. "
                 f"Please make sure to run `prepare_model` before passing it in to the "
                 f"Analyzer."
             )
             self.logger.error(error_msg)
-            raise UnsupportableModuleError(error_msg)
+            raise UnsupportableModuleError(error_msg) from e
         self.logger.info(f"Tracking modules with names: {tracked_module_names}.")
 
         if self.state.use_distributed and not isinstance(model, (DDP, FSDP)):