finalize

neuralmagic · Feb 13, 2023 · a6a3764 · a6a3764
1 parent a8a6992
commit a6a3764
Showing 1 changed file with 7 additions and 0 deletions.
diff --git a/src/sparseml/pytorch/utils/logger.py b/src/sparseml/pytorch/utils/logger.py
@@ -534,7 +534,14 @@ def __init__(
         init_kwargs: Optional[Dict] = None,
         name: str = "wandb",
         enabled: bool = True,
+        wandb_err: Optional[Exception] = wandb_err,
     ):
+        if wandb_err:
+            raise ModuleNotFoundError(
+                "Error: Failed to import wandb. "
+                "Please install the wandb library in order to use it."
+            ) from wandb_err
+
         super().__init__(
             lambda_func=self._log_lambda,
             name=name,