better solution

src/sparseml/pytorch/utils/logger.py

@@ -513,18 +513,6 @@ class WANDBLogger(LambdaLogger):
     :param enabled: True to log, False otherwise
     """
 
-    try:
-        import wandb
-
-        wandb_err = None
-    except Exception as err:
-        wandb = None
-        wandb_err = err
-        raise ModuleNotFoundError(
-            "Error: Failed to import wandb. "
-            "Please install the wandb library in order to use it."
-        ) from wandb_err
-
     def available(self) -> bool:
         """
         :return: True if wandb is available and installed, False, otherwise
@@ -537,6 +525,21 @@ def __init__(
         name: str = "wandb",
         enabled: bool = True,
     ):
+        try:
+            import wandb
+
+            wandb_err = None
+        except Exception as err:
+            wandb = None
+            wandb_err = err
+            raise ModuleNotFoundError(
+                "Error: Failed to import wandb. "
+                "Please install the wandb library in order to use it."
+            ) from wandb_err
+
+        self.wandb = wandb
+        self.wandb_err = wandb_err
+
         super().__init__(
             lambda_func=self._log_lambda,
             name=name,