nshtrainer 0.25.0__tar.gz → 0.26.0__tar.gz

{nshtrainer-0.25.0 → nshtrainer-0.26.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: nshtrainer
-Version: 0.25.0
+Version: 0.26.0
 Summary: 
 Author: Nima Shoghi
 Author-email: nimashoghi@gmail.com
@@ -22,7 +22,7 @@ Requires-Dist: psutil
 Requires-Dist: pytorch-lightning
 Requires-Dist: tensorboard ; extra == "extra"
 Requires-Dist: torch
-Requires-Dist: torchmetrics ; extra == "extra"
+Requires-Dist: torchmetrics
 Requires-Dist: typing-extensions
 Requires-Dist: wandb ; extra == "extra"
 Requires-Dist: wrapt ; extra == "extra"

{nshtrainer-0.25.0 → nshtrainer-0.26.0}/pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "nshtrainer"
-version = "0.25.0"
+version = "0.26.0"
 description = ""
 authors = ["Nima Shoghi <nimashoghi@gmail.com>"]
 readme = "README.md"
@@ -17,7 +17,7 @@ typing-extensions = "*"
 packaging = "*"
 lightning = "*"
 pytorch-lightning = "*"
-torchmetrics = { version = "*", optional = true }
+torchmetrics = "*"
 wrapt = { version = "*", optional = true }
 GitPython = { version = "*", optional = true }
 wandb = { version = "*", optional = true }
@@ -46,11 +46,4 @@ reportPrivateImportUsage = false
 ignore = ["F722", "F821", "E731", "E741"]
 
 [tool.poetry.extras]
-extra = [
-    "torchmetrics",
-    "wrapt",
-    "GitPython",
-    "wandb",
-    "tensorboard",
-    "huggingface-hub",
-]
+extra = ["wrapt", "GitPython", "wandb", "tensorboard", "huggingface-hub"]

{nshtrainer-0.25.0 → nshtrainer-0.26.0}/src/nshtrainer/callbacks/gradient_skipping.py

@@ -1,8 +1,8 @@
-import importlib.util
 import logging
 from typing import Any, Literal, Protocol, runtime_checkable
 
 import torch
+import torchmetrics
 from lightning.pytorch import Callback, LightningModule, Trainer
 from torch.optim import Optimizer
 from typing_extensions import override
@@ -20,19 +20,12 @@ class HasGradSkippedSteps(Protocol):
 
 class GradientSkipping(Callback):
     def __init__(self, config: "GradientSkippingConfig"):
-        if importlib.util.find_spec("torchmetrics") is not None:
-            raise ImportError(
-                "To use the GradientSkipping callback, please install torchmetrics: pip install torchmetrics"
-            )
-
         super().__init__()
         self.config = config
 
     @override
     def setup(self, trainer: Trainer, pl_module: LightningModule, stage: str) -> None:
         if not isinstance(pl_module, HasGradSkippedSteps):
-            import torchmetrics  # type: ignore
-
             pl_module.grad_skipped_steps = torchmetrics.SumMetric()
 
     @override