move export to class

src/sparseml/pytorch/utils/logger.py

@@ -36,19 +36,6 @@
     SummaryWriter = object
     tensorboard_import_error = tensorboard_err
 
-
-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
-
 from sparseml.utils import ALL_TOKEN, create_dirs
 
 
@@ -526,12 +513,23 @@ class WANDBLogger(LambdaLogger):
     :param enabled: True to log, False otherwise
     """
 
-    @staticmethod
-    def available() -> bool:
+    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
         """
-        return not wandb_err
+        return not self.wandb_err
 
     def __init__(
         self,
@@ -553,15 +551,15 @@ def __init__(
             else:
                 init_kwargs = {"dir": test_log_path}
 
-        if wandb_err:
-            raise wandb_err
+        if self.wandb_err:
+            raise self.wandb_err
 
         if init_kwargs:
-            wandb.init(**init_kwargs)
+            self.wandb.init(**init_kwargs)
         else:
-            wandb.init()
+            self.wandb.init()
 
-        self.wandb = wandb
+        self.wandb = self.wandb
 
     def _log_lambda(
         self,
@@ -582,7 +580,7 @@ def _log_lambda(
                 values = {f"{tag}/{key}": val for key, val in values.items()}
             params.update(values)
 
-        wandb.log(params, step=step)
+        self.wandb.log(params, step=step)
 
         return True
 
@@ -593,7 +591,7 @@ def save(
         """
         :param file_path: path to a file to be saved
         """
-        wandb.save(file_path)
+        self.wandb.save(file_path)
         return True