From ed9673cead48d1aa7c619ec666b851fbc68973af Mon Sep 17 00:00:00 2001
From: dbogunowicz <97082108+dbogunowicz@users.noreply.github.com>
Date: Mon, 13 Feb 2023 16:28:12 +0100
Subject: [PATCH] [Cherry Pick] [Fix] Add more verbose`wandb` import error
 (#1375)

* clean cherry pick

* new solution

* final fix
---
 src/sparseml/pytorch/utils/logger.py | 7 +++++++
 1 file changed, 7 insertions(+)

diff --git a/src/sparseml/pytorch/utils/logger.py b/src/sparseml/pytorch/utils/logger.py
index 9d7cc4e3c2d..66ea8e920ac 100644
--- 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,