tey

Lightning-Universe · Jun 30, 2023 · 031fad5 · 031fad5
1 parent 676aa88
commit 031fad5
Show file tree

Hide file tree

Showing 2 changed files with 5 additions and 8 deletions.
diff --git a/src/pl_bolts/callbacks/sparseml.py b/src/pl_bolts/callbacks/sparseml.py
@@ -14,19 +14,11 @@
 from typing import Any, Optional
 
 import torch
-from lightning_utilities.core.rank_zero import rank_zero_warn
 from pytorch_lightning import Callback, LightningModule, Trainer
 from pytorch_lightning.utilities.exceptions import MisconfigurationException
 
 from pl_bolts.utils import _SPARSEML_AVAILABLE
 
-if _SPARSEML_AVAILABLE:
-    try:
-        import sparseml
-    except ImportError as err:
-        _SPARSEML_AVAILABLE = False
-        rank_zero_warn(str(err))
-
 if _SPARSEML_AVAILABLE:
     from sparseml.pytorch.optim import ScheduledModifierManager
     from sparseml.pytorch.utils import ModuleExporter

diff --git a/src/pl_bolts/utils/__init__.py b/src/pl_bolts/utils/__init__.py
@@ -26,5 +26,10 @@
 _MATPLOTLIB_AVAILABLE: bool = module_available("matplotlib")
 _JSONARGPARSE_GREATER_THAN_4_16_0 = compare_version("jsonargparse", operator.gt, "4.16.0")
 
+if _SPARSEML_AVAILABLE:
+    import sparseml
+
+    _SPARSEML_AVAILABLE = compare_version("sparseml", operator.le, sparseml.pytorch.base._TORCH_MAX_VERSION)
+
 
 __all__ = ["BatchGradientVerification"]