nshtrainer 0.13.1__py3-none-any.whl → 0.14.1__py3-none-any.whl

nshtrainer/callbacks/early_stopping.py

@@ -73,9 +73,7 @@ class EarlyStopping(_EarlyStopping):
 
     @override
     @staticmethod
-    def _log_info(
-        trainer: Trainer | None, message: str, log_rank_zero_only: bool
-    ) -> None:
+    def _log_info(trainer: Trainer | None, message: str, log_rank_zero_only: bool):
         rank = _get_rank()
         if trainer is not None and trainer.world_size <= 1:
             rank = None

nshtrainer/ll/__init__.py

@@ -23,7 +23,6 @@ from .log import pretty as pretty
 from .lr_scheduler import LRSchedulerConfig as LRSchedulerConfig
 from .model import Base as Base
 from .model import BaseConfig as BaseConfig
-from .model import BaseLoggerConfig as BaseLoggerConfig
 from .model import BaseProfilerConfig as BaseProfilerConfig
 from .model import CheckpointLoadingConfig as CheckpointLoadingConfig
 from .model import CheckpointSavingConfig as CheckpointSavingConfig
@@ -48,7 +47,6 @@ from .model import PrimaryMetricConfig as PrimaryMetricConfig
 from .model import ReproducibilityConfig as ReproducibilityConfig
 from .model import SanityCheckingConfig as SanityCheckingConfig
 from .model import TrainerConfig as TrainerConfig
-from .model import WandbWatchConfig as WandbWatchConfig
 from .nn import TypedModuleDict as TypedModuleDict
 from .nn import TypedModuleList as TypedModuleList
 from .optimizer import OptimizerConfig as OptimizerConfig

nshtrainer/loggers/__init__.py

@@ -0,0 +1,13 @@
+from typing import Annotated, TypeAlias
+
+import nshconfig as C
+
+from ._base import BaseLoggerConfig as BaseLoggerConfig
+from .csv import CSVLoggerConfig as CSVLoggerConfig
+from .tensorboard import TensorboardLoggerConfig as TensorboardLoggerConfig
+from .wandb import WandbLoggerConfig as WandbLoggerConfig
+
+LoggerConfig: TypeAlias = Annotated[
+    CSVLoggerConfig | TensorboardLoggerConfig | WandbLoggerConfig,
+    C.Field(discriminator="name"),
+]

nshtrainer/loggers/_base.py

@@ -0,0 +1,26 @@
+from abc import ABC, abstractmethod
+from typing import TYPE_CHECKING
+
+import nshconfig as C
+from lightning.pytorch.loggers import Logger
+
+if TYPE_CHECKING:
+    from ..model import BaseConfig
+
+
+class BaseLoggerConfig(C.Config, ABC):
+    enabled: bool = True
+    """Enable this logger."""
+
+    priority: int = 0
+    """Priority of the logger. Higher priority loggers are created first."""
+
+    log_dir: C.DirectoryPath | None = None
+    """Directory to save the logs to. If None, will use the default log directory for the trainer."""
+
+    @abstractmethod
+    def create_logger(self, root_config: "BaseConfig") -> Logger | None: ...
+
+    def disable_(self):
+        self.enabled = False
+        return self

nshtrainer/loggers/csv.py

@@ -0,0 +1,40 @@
+from typing import Literal
+
+from typing_extensions import override
+
+from ._base import BaseLoggerConfig
+
+
+class CSVLoggerConfig(BaseLoggerConfig):
+    name: Literal["csv"] = "csv"
+
+    enabled: bool = True
+    """Enable CSV logging."""
+
+    priority: int = 0
+    """Priority of the logger. Higher priority loggers are created first."""
+
+    prefix: str = ""
+    """A string to put at the beginning of metric keys."""
+
+    flush_logs_every_n_steps: int = 100
+    """How often to flush logs to disk."""
+
+    @override
+    def create_logger(self, root_config):
+        if not self.enabled:
+            return None
+
+        from lightning.pytorch.loggers.csv_logs import CSVLogger
+
+        save_dir = root_config.directory._resolve_log_directory_for_logger(
+            root_config.id,
+            self,
+        )
+        return CSVLogger(
+            save_dir=save_dir,
+            name=root_config.run_name,
+            version=root_config.id,
+            prefix=self.prefix,
+            flush_logs_every_n_steps=self.flush_logs_every_n_steps,
+        )

nshtrainer/loggers/tensorboard.py

@@ -0,0 +1,73 @@
+import logging
+from typing import Literal
+
+import nshconfig as C
+from typing_extensions import override
+
+from ._base import BaseLoggerConfig
+
+log = logging.getLogger(__name__)
+
+
+def _tensorboard_available():
+    try:
+        from lightning.fabric.loggers.tensorboard import (
+            _TENSORBOARD_AVAILABLE,
+            _TENSORBOARDX_AVAILABLE,
+        )
+
+        if not _TENSORBOARD_AVAILABLE and not _TENSORBOARDX_AVAILABLE:
+            log.warning(
+                "TensorBoard/TensorBoardX not found. Disabling TensorBoardLogger. "
+                "Please install TensorBoard with `pip install tensorboard` or "
+                "TensorBoardX with `pip install tensorboardx` to enable TensorBoard logging."
+            )
+            return False
+        return True
+    except ImportError:
+        return False
+
+
+class TensorboardLoggerConfig(BaseLoggerConfig):
+    name: Literal["tensorboard"] = "tensorboard"
+
+    enabled: bool = C.Field(default_factory=lambda: _tensorboard_available())
+    """Enable TensorBoard logging."""
+
+    priority: int = 2
+    """Priority of the logger. Higher priority loggers are created first."""
+
+    log_graph: bool = False
+    """
+    Adds the computational graph to tensorboard. This requires that
+        the user has defined the `self.example_input_array` attribute in their
+        model.
+    """
+
+    default_hp_metric: bool = True
+    """
+    Enables a placeholder metric with key `hp_metric` when `log_hyperparams` is
+        called without a metric (otherwise calls to log_hyperparams without a metric are ignored).
+    """
+
+    prefix: str = ""
+    """A string to put at the beginning of metric keys."""
+
+    @override
+    def create_logger(self, root_config):
+        if not self.enabled:
+            return None
+
+        from lightning.pytorch.loggers.tensorboard import TensorBoardLogger
+
+        save_dir = root_config.directory._resolve_log_directory_for_logger(
+            root_config.id,
+            self,
+        )
+        return TensorBoardLogger(
+            save_dir=save_dir,
+            name=root_config.run_name,
+            version=root_config.id,
+            log_graph=self.log_graph,
+            default_hp_metric=self.default_hp_metric,
+        )

nshtrainer/loggers/wandb.py

@@ -0,0 +1,163 @@
+import logging
+from typing import TYPE_CHECKING, Literal
+
+import nshconfig as C
+import pkg_resources
+from lightning.pytorch import Callback, LightningModule, Trainer
+from typing_extensions import override
+
+from ..callbacks import WandbWatchConfig
+from ..callbacks.base import CallbackConfigBase
+from ._base import BaseLoggerConfig
+
+if TYPE_CHECKING:
+    from ..model.config import BaseConfig
+
+log = logging.getLogger(__name__)
+
+
+def _project_name(
+    root_config: "BaseConfig",
+    default_project: str = "lightning_logs",
+):
+    # If the config has a project name, use that.
+    if project := root_config.project:
+        return project
+
+    # Otherwise, we should use the name of the module that the config is defined in,
+    #   if we can find it.
+    # If this isn't in a module, use the default project name.
+    if not (module := root_config.__module__):
+        return default_project
+
+    # If the module is a package, use the package name.
+    if not (module := module.split(".", maxsplit=1)[0].strip()):
+        return default_project
+
+    return module
+
+
+def _wandb_available():
+    try:
+        from lightning.pytorch.loggers.wandb import _WANDB_AVAILABLE
+
+        if not _WANDB_AVAILABLE:
+            log.warning("WandB not found. Disabling WandbLogger.")
+            return False
+        return True
+    except ImportError:
+        return False
+
+
+class FinishWandbOnTeardownCallback(Callback):
+    @override
+    def teardown(
+        self,
+        trainer: Trainer,
+        pl_module: LightningModule,
+        stage: str,
+    ):
+        try:
+            import wandb  # type: ignore
+        except ImportError:
+            return
+
+        if wandb.run is None:
+            return
+
+        wandb.finish()
+
+
+class WandbLoggerConfig(CallbackConfigBase, BaseLoggerConfig):
+    name: Literal["wandb"] = "wandb"
+
+    enabled: bool = C.Field(default_factory=lambda: _wandb_available())
+    """Enable WandB logging."""
+
+    priority: int = 2
+    """Priority of the logger. Higher priority loggers are created first,
+    and the highest priority logger is the "main" logger for PyTorch Lightning."""
+
+    project: str | None = None
+    """WandB project name to use for the logger. If None, will use the root config's project name."""
+
+    log_model: bool | Literal["all"] = False
+    """
+    Whether to log the model checkpoints to wandb.
+    Valid values are:
+        - False: Do not log the model checkpoints.
+        - True: Log the latest model checkpoint.
+        - "all": Log all model checkpoints.
+    """
+
+    watch: WandbWatchConfig | None = WandbWatchConfig()
+    """WandB model watch configuration. Used to log model architecture, gradients, and parameters."""
+
+    offline: bool = False
+    """Whether to run WandB in offline mode."""
+
+    use_wandb_core: bool = True
+    """Whether to use the new `wandb-core` backend for WandB.
+    `wandb-core` is a new backend for WandB that is faster and more efficient than the old backend.
+    """
+
+    def offline_(self, value: bool = True):
+        self.offline = value
+        return self
+
+    def core_(self, value: bool = True):
+        self.use_wandb_core = value
+        return self
+
+    @override
+    def create_logger(self, root_config):
+        if not self.enabled:
+            return None
+
+        # If `wandb-core` is enabled, we should use the new backend.
+        if self.use_wandb_core:
+            try:
+                import wandb  # type: ignore
+
+                # The minimum version that supports the new backend is 0.17.5
+                if pkg_resources.parse_version(
+                    wandb.__version__
+                ) < pkg_resources.parse_version("0.17.5"):
+                    log.warning(
+                        "The version of WandB installed does not support the `wandb-core` backend "
+                        f"(expected version >= 0.17.5, found version {wandb.__version__}). "
+                        "Please either upgrade to a newer version of WandB or disable the `use_wandb_core` option."
+                    )
+                else:
+                    wandb.require("core")
+                    log.critical("Using the `wandb-core` backend for WandB.")
+            except ImportError:
+                pass
+
+        from lightning.pytorch.loggers.wandb import WandbLogger
+
+        save_dir = root_config.directory._resolve_log_directory_for_logger(
+            root_config.id,
+            self,
+        )
+        return WandbLogger(
+            save_dir=save_dir,
+            project=self.project or _project_name(root_config),
+            name=root_config.run_name,
+            version=root_config.id,
+            log_model=self.log_model,
+            notes=(
+                "\n".join(f"- {note}" for note in root_config.notes)
+                if root_config.notes
+                else None
+            ),
+            tags=root_config.tags,
+            offline=self.offline,
+        )
+
+    @override
+    def create_callbacks(self, root_config):
+        yield FinishWandbOnTeardownCallback()
+
+        if self.watch:
+            yield from self.watch.create_callbacks(root_config)

nshtrainer/model/__init__.py

@@ -3,7 +3,6 @@ from typing_extensions import TypeAlias
 from .base import Base as Base
 from .base import LightningModuleBase as LightningModuleBase
 from .config import BaseConfig as BaseConfig
-from .config import BaseLoggerConfig as BaseLoggerConfig
 from .config import BaseProfilerConfig as BaseProfilerConfig
 from .config import BestCheckpointCallbackConfig as BestCheckpointCallbackConfig
 from .config import CheckpointLoadingConfig as CheckpointLoadingConfig
@@ -22,6 +21,5 @@ from .config import PrimaryMetricConfig as PrimaryMetricConfig
 from .config import ReproducibilityConfig as ReproducibilityConfig
 from .config import SanityCheckingConfig as SanityCheckingConfig
 from .config import TrainerConfig as TrainerConfig
-from .config import WandbWatchConfig as WandbWatchConfig
 
 ConfigList: TypeAlias = list[tuple[BaseConfig, type[LightningModuleBase]]]

nshtrainer/model/config.py

@@ -19,7 +19,6 @@ from typing import (
 
 import nshconfig as C
 import numpy as np
-import pkg_resources
 import torch
 from lightning.fabric.plugins import CheckpointIO, ClusterEnvironment
 from lightning.fabric.plugins.precision.precision import _PRECISION_INPUT
@@ -31,7 +30,6 @@ from lightning.pytorch.plugins.layer_sync import LayerSync
 from lightning.pytorch.plugins.precision.precision import Precision
 from lightning.pytorch.profilers import Profiler
 from lightning.pytorch.strategies.strategy import Strategy
-from pydantic import DirectoryPath
 from typing_extensions import Self, TypedDict, TypeVar, override
 
 from .._checkpoint.loader import CheckpointLoadingConfig
@@ -41,9 +39,14 @@ from ..callbacks import (
     EarlyStoppingConfig,
     LastCheckpointCallbackConfig,
     OnExceptionCheckpointCallbackConfig,
-    WandbWatchConfig,
 )
 from ..callbacks.base import CallbackConfigBase
+from ..loggers import (
+    CSVLoggerConfig,
+    LoggerConfig,
+    TensorboardLoggerConfig,
+    WandbLoggerConfig,
+)
 from ..metrics import MetricConfig
 from ..util._environment_info import EnvironmentConfig
 
@@ -205,255 +208,6 @@ ProfilerConfig: TypeAlias = Annotated[
 ]
 
 
-class BaseLoggerConfig(C.Config, ABC):
-    enabled: bool = True
-    """Enable this logger."""
-
-    priority: int = 0
-    """Priority of the logger. Higher priority loggers are created first."""
-
-    log_dir: DirectoryPath | None = None
-    """Directory to save the logs to. If None, will use the default log directory for the trainer."""
-
-    @abstractmethod
-    def create_logger(self, root_config: "BaseConfig") -> Logger | None: ...
-
-    def disable_(self):
-        self.enabled = False
-        return self
-
-
-def _project_name(
-    root_config: "BaseConfig",
-    default_project: str = "lightning_logs",
-):
-    # If the config has a project name, use that.
-    if project := root_config.project:
-        return project
-
-    # Otherwise, we should use the name of the module that the config is defined in,
-    #   if we can find it.
-    # If this isn't in a module, use the default project name.
-    if not (module := root_config.__module__):
-        return default_project
-
-    # If the module is a package, use the package name.
-    if not (module := module.split(".", maxsplit=1)[0].strip()):
-        return default_project
-
-    return module
-
-
-def _wandb_available():
-    try:
-        from lightning.pytorch.loggers.wandb import _WANDB_AVAILABLE
-
-        if not _WANDB_AVAILABLE:
-            log.warning("WandB not found. Disabling WandbLogger.")
-            return False
-        return True
-    except ImportError:
-        return False
-
-
-class WandbLoggerConfig(CallbackConfigBase, BaseLoggerConfig):
-    name: Literal["wandb"] = "wandb"
-
-    enabled: bool = C.Field(default_factory=lambda: _wandb_available())
-    """Enable WandB logging."""
-
-    priority: int = 2
-    """Priority of the logger. Higher priority loggers are created first,
-    and the highest priority logger is the "main" logger for PyTorch Lightning."""
-
-    project: str | None = None
-    """WandB project name to use for the logger. If None, will use the root config's project name."""
-
-    log_model: bool | Literal["all"] = False
-    """
-    Whether to log the model checkpoints to wandb.
-    Valid values are:
-        - False: Do not log the model checkpoints.
-        - True: Log the latest model checkpoint.
-        - "all": Log all model checkpoints.
-    """
-
-    watch: WandbWatchConfig = WandbWatchConfig()
-    """WandB model watch configuration. Used to log model architecture, gradients, and parameters."""
-
-    offline: bool = False
-    """Whether to run WandB in offline mode."""
-
-    use_wandb_core: bool = True
-    """Whether to use the new `wandb-core` backend for WandB.
-    `wandb-core` is a new backend for WandB that is faster and more efficient than the old backend.
-    """
-
-    def offline_(self, value: bool = True):
-        self.offline = value
-        return self
-
-    def core_(self, value: bool = True):
-        self.use_wandb_core = value
-        return self
-
-    @override
-    def create_logger(self, root_config):
-        if not self.enabled:
-            return None
-
-        # If `wandb-core` is enabled, we should use the new backend.
-        if self.use_wandb_core:
-            try:
-                import wandb  # type: ignore
-
-                # The minimum version that supports the new backend is 0.17.5
-                if pkg_resources.parse_version(
-                    wandb.__version__
-                ) < pkg_resources.parse_version("0.17.5"):
-                    log.warning(
-                        "The version of WandB installed does not support the `wandb-core` backend "
-                        f"(expected version >= 0.17.5, found version {wandb.__version__}). "
-                        "Please either upgrade to a newer version of WandB or disable the `use_wandb_core` option."
-                    )
-                else:
-                    wandb.require("core")
-                    log.critical("Using the `wandb-core` backend for WandB.")
-            except ImportError:
-                pass
-
-        from lightning.pytorch.loggers.wandb import WandbLogger
-
-        save_dir = root_config.directory._resolve_log_directory_for_logger(
-            root_config.id,
-            self,
-        )
-        return WandbLogger(
-            save_dir=save_dir,
-            project=self.project or _project_name(root_config),
-            name=root_config.run_name,
-            version=root_config.id,
-            log_model=self.log_model,
-            notes=(
-                "\n".join(f"- {note}" for note in root_config.notes)
-                if root_config.notes
-                else None
-            ),
-            tags=root_config.tags,
-            offline=self.offline,
-        )
-
-    @override
-    def create_callbacks(self, root_config):
-        if self.watch:
-            yield from self.watch.create_callbacks(root_config)
-
-
-class CSVLoggerConfig(BaseLoggerConfig):
-    name: Literal["csv"] = "csv"
-
-    enabled: bool = True
-    """Enable CSV logging."""
-
-    priority: int = 0
-    """Priority of the logger. Higher priority loggers are created first."""
-
-    prefix: str = ""
-    """A string to put at the beginning of metric keys."""
-
-    flush_logs_every_n_steps: int = 100
-    """How often to flush logs to disk."""
-
-    @override
-    def create_logger(self, root_config):
-        if not self.enabled:
-            return None
-
-        from lightning.pytorch.loggers.csv_logs import CSVLogger
-
-        save_dir = root_config.directory._resolve_log_directory_for_logger(
-            root_config.id,
-            self,
-        )
-        return CSVLogger(
-            save_dir=save_dir,
-            name=root_config.run_name,
-            version=root_config.id,
-            prefix=self.prefix,
-            flush_logs_every_n_steps=self.flush_logs_every_n_steps,
-        )
-
-
-def _tensorboard_available():
-    try:
-        from lightning.fabric.loggers.tensorboard import (
-            _TENSORBOARD_AVAILABLE,
-            _TENSORBOARDX_AVAILABLE,
-        )
-
-        if not _TENSORBOARD_AVAILABLE and not _TENSORBOARDX_AVAILABLE:
-            log.warning(
-                "TensorBoard/TensorBoardX not found. Disabling TensorBoardLogger. "
-                "Please install TensorBoard with `pip install tensorboard` or "
-                "TensorBoardX with `pip install tensorboardx` to enable TensorBoard logging."
-            )
-            return False
-        return True
-    except ImportError:
-        return False
-
-
-class TensorboardLoggerConfig(BaseLoggerConfig):
-    name: Literal["tensorboard"] = "tensorboard"
-
-    enabled: bool = C.Field(default_factory=lambda: _tensorboard_available())
-    """Enable TensorBoard logging."""
-
-    priority: int = 2
-    """Priority of the logger. Higher priority loggers are created first."""
-
-    log_graph: bool = False
-    """
-    Adds the computational graph to tensorboard. This requires that
-        the user has defined the `self.example_input_array` attribute in their
-        model.
-    """
-
-    default_hp_metric: bool = True
-    """
-    Enables a placeholder metric with key `hp_metric` when `log_hyperparams` is
-        called without a metric (otherwise calls to log_hyperparams without a metric are ignored).
-    """
-
-    prefix: str = ""
-    """A string to put at the beginning of metric keys."""
-
-    @override
-    def create_logger(self, root_config):
-        if not self.enabled:
-            return None
-
-        from lightning.pytorch.loggers.tensorboard import TensorBoardLogger
-
-        save_dir = root_config.directory._resolve_log_directory_for_logger(
-            root_config.id,
-            self,
-        )
-        return TensorBoardLogger(
-            save_dir=save_dir,
-            name=root_config.run_name,
-            version=root_config.id,
-            log_graph=self.log_graph,
-            default_hp_metric=self.default_hp_metric,
-        )
-
-
-LoggerConfig: TypeAlias = Annotated[
-    WandbLoggerConfig | CSVLoggerConfig | TensorboardLoggerConfig,
-    C.Field(discriminator="name"),
-]
-
-
 class LoggingConfig(CallbackConfigBase):
     enabled: bool = True
     """Enable experiment tracking."""
@@ -474,29 +228,32 @@ class LoggingConfig(CallbackConfigBase):
     """If enabled, will automatically save logged metrics using ActSave (if nshutils is installed)."""
 
     @property
-    def wandb(self) -> WandbLoggerConfig | None:
+    def wandb(self):
         return next(
             (
                 logger
                 for logger in self.loggers
                 if isinstance(logger, WandbLoggerConfig)
             ),
+            None,
         )
 
     @property
-    def csv(self) -> CSVLoggerConfig | None:
+    def csv(self):
         return next(
             (logger for logger in self.loggers if isinstance(logger, CSVLoggerConfig)),
+            None,
         )
 
     @property
-    def tensorboard(self) -> TensorboardLoggerConfig | None:
+    def tensorboard(self):
         return next(
             (
                 logger
                 for logger in self.loggers
                 if isinstance(logger, TensorboardLoggerConfig)
             ),
+            None,
         )
 
     def create_loggers(self, root_config: "BaseConfig"):
@@ -509,9 +266,8 @@ class LoggingConfig(CallbackConfigBase):
         Returns:
             list[Logger]: A list of constructed loggers.
         """
-        loggers: list[Logger] = []
         if not self.enabled:
-            return loggers
+            return
 
         for logger_config in sorted(
             self.loggers,
@@ -522,8 +278,7 @@ class LoggingConfig(CallbackConfigBase):
                 continue
             if (logger := logger_config.create_logger(root_config)) is None:
                 continue
-            loggers.append(logger)
-        return loggers
+            yield logger
 
     @override
     def create_callbacks(self, root_config):

nshtrainer/trainer/trainer.py

@@ -244,7 +244,7 @@ class Trainer(LightningTrainer):
             log.critical(f"Disabling logger because {config.trainer.logging.enabled=}.")
             kwargs["logger"] = False
         else:
-            _update_kwargs(logger=config.trainer.logging.create_loggers(config))
+            _update_kwargs(logger=list(config.trainer.logging.create_loggers(config)))
 
         if config.trainer.auto_determine_num_nodes:
             # When num_nodes is auto, we need to detect the number of nodes.

{nshtrainer-0.13.1.dist-info → nshtrainer-0.14.1.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: nshtrainer
-Version: 0.13.1
+Version: 0.14.1
 Summary: 
 Author: Nima Shoghi
 Author-email: nimashoghi@gmail.com

{nshtrainer-0.13.1.dist-info → nshtrainer-0.14.1.dist-info}/RECORD

@@ -15,7 +15,7 @@ nshtrainer/callbacks/checkpoint/_base.py,sha256=yBUZ63P4a3m5jviIUGxBd2WXq_TigpKt
 nshtrainer/callbacks/checkpoint/best_checkpoint.py,sha256=DJiLo7NDzd-lp-O3v7Cv8WejyjXPV_6_RmfltKO9fvE,2165
 nshtrainer/callbacks/checkpoint/last_checkpoint.py,sha256=CqB_8Xv32rtpLCaEEPi6DbRZm4ph5TWS-LfqIHXUIUA,1097
 nshtrainer/callbacks/checkpoint/on_exception_checkpoint.py,sha256=ctT88EGT22_t_6tr5r7Sfo43cuve6XeroBnBYRMPOus,3372
-nshtrainer/callbacks/early_stopping.py,sha256=DlB3FOMbPYaCn1pYHiqa_upnoBExAMAbF2oHA92qMUs,4639
+nshtrainer/callbacks/early_stopping.py,sha256=VWuJz0oN87b6SwBeVc32YNpeJr1wts8K45k8JJJmG9I,4617
 nshtrainer/callbacks/ema.py,sha256=8-WHmKFP3VfnzMviJaIFmVD9xHPqIPmq9NRF5xdu3c8,12131
 nshtrainer/callbacks/finite_checks.py,sha256=gJC_RUr3ais3FJI0uB6wUZnDdE3WRwCix3ppA3PwQXA,2077
 nshtrainer/callbacks/gradient_skipping.py,sha256=pqu5AELx4ctJxR2Y7YSSiGd5oGauVCTZFCEIIS6s88w,3665
@@ -29,7 +29,7 @@ nshtrainer/callbacks/wandb_watch.py,sha256=Y6SEXfIx3kDDQbI5zpP53BVq0FBLJbLd3RJsi
 nshtrainer/data/__init__.py,sha256=7mk1tr7SWUZ7ySbsf0y0ZPszk7u4QznPhQ-7wnpH9ec,149
 nshtrainer/data/balanced_batch_sampler.py,sha256=dGBTDDtlBU6c-ZlVQOCnTW7SjTB5hczWsOWEdUWjvkA,4385
 nshtrainer/data/transform.py,sha256=6SNs3_TpNpfhcwTwvPKyEJ3opM1OT7LmMEYQNHKgRl8,2227
-nshtrainer/ll/__init__.py,sha256=dD0ISxHJ2lg1HLSM0b3db7TBlsPpQCtChnuYO-c2oqI,2635
+nshtrainer/ll/__init__.py,sha256=6UTt2apSD8tOZw3M7hyd-33v4RKSpNNATlWFbW4cNnU,2523
 nshtrainer/ll/_experimental.py,sha256=oBQCKOEVYoxuUU9eLb-Fg2B2mzZD7SA0zfAO6lmWZ88,53
 nshtrainer/ll/actsave.py,sha256=2lbiseSrjcwFT6AiyLNWarTWl1bnzliVWlu1iOfnP30,209
 nshtrainer/ll/callbacks.py,sha256=AxyUmc8aGRSjx6WwwgXYCmdJ73rwLuEAEH0AGRosojQ,49
@@ -46,15 +46,20 @@ nshtrainer/ll/snoop.py,sha256=hG9VCdm8mIZytHLZgUKRSoWqg55rVvUpBAH-OSiwgvI,36
 nshtrainer/ll/trainer.py,sha256=hkn2xPtrSPQ7LqQhbyAKuMfNyHdhqB9bDPvgRCK8oJM,47
 nshtrainer/ll/typecheck.py,sha256=ryV1Tzcf7hJ4I19H1oQVkikU9spmRk8jyIKQZ5UF7pQ,62
 nshtrainer/ll/util.py,sha256=PQGu5Ff1raizxjXdm2rFFu4Mo816dBIYmItjkJX1_qk,44
+nshtrainer/loggers/__init__.py,sha256=C_xk0A3_qKbNdTmzK85AgjRHFD3w-jPRS2ig-iPhfEk,448
+nshtrainer/loggers/_base.py,sha256=xiZKEK0ALJkcqf4OpVNRY0QbZsamR_WR7x7m_68YHXQ,705
+nshtrainer/loggers/csv.py,sha256=D_lYyd94bZ8jAgnRo-ARtFgVcInaD9zktxtsUD9RWCI,1052
+nshtrainer/loggers/tensorboard.py,sha256=wL2amRSdP68zbslZvBeM0ZQBnjF3hIKsz-_lBbdomaM,2216
+nshtrainer/loggers/wandb.py,sha256=QteSPfaBWxQ-FuuaF6_rKxw8twJdg_6oHoTQlbqvrXk,5079
 nshtrainer/lr_scheduler/__init__.py,sha256=uEvgaFAs-4s_bAEMaildy0GT6OvgpgOEKTuzqutESHE,736
 nshtrainer/lr_scheduler/_base.py,sha256=7xOIuxQ86YHbFWG5a3gX46emQj1WN_LaY4-i0Q1TDBg,3659
 nshtrainer/lr_scheduler/linear_warmup_cosine.py,sha256=mn6cyizyI_stkXtg6zxIEGF9btIxMRWigUHUTlUYCSw,5221
 nshtrainer/lr_scheduler/reduce_lr_on_plateau.py,sha256=h76oTHYpMxauV_l6lviya5DW-WKArwxxf7ZQizhmbCw,2782
 nshtrainer/metrics/__init__.py,sha256=ObLIELGguIEcUpRsUkqh1ltrvZii6vglTpJGrPvoy00,50
 nshtrainer/metrics/_config.py,sha256=jgRBfDAQLFTW7AiUY7CRtdfts6CR6keeuqm0FFMWCzQ,1288
-nshtrainer/model/__init__.py,sha256=RlGW5a46DZcqK6cYICYxDaKpZIEj-8zLxoMrl432tno,1429
+nshtrainer/model/__init__.py,sha256=BmqSbf6v6oyeilti4iEn_Tyrr1kRmcFcJekTb8NeglI,1315
 nshtrainer/model/base.py,sha256=AXRfEsFAT0Ln7zjYVPU5NgtHS_c8FZM-M4pyLamO7OA,17516
-nshtrainer/model/config.py,sha256=gvZRv8Gow6ViF0t8OIoR_IzxLn8GzVmbpvoVzpdykFc,50871
+nshtrainer/model/config.py,sha256=D6Y-Y7GoMrpo7A2dmIqJsqc4X2IHwyl9OEHxO4uOc0g,42918
 nshtrainer/model/modules/callback.py,sha256=K0-cyEtBcQhI7Q2e-AGTE8T-GghUPY9DYmneU6ULV6g,6401
 nshtrainer/model/modules/debug.py,sha256=Yy7XEdPou9BkCsD5hJchwJGmCVGrfUru5g9VjPM4uAw,1120
 nshtrainer/model/modules/distributed.py,sha256=ABpR9d-3uBS_fivfy_WYW-dExW6vp5BPaoPQnOudHng,1725
@@ -74,7 +79,7 @@ nshtrainer/trainer/__init__.py,sha256=P2rmr8oBVTHk-HJHYPcUwWqDEArMbPR4_rPpATbWK3
 nshtrainer/trainer/_runtime_callback.py,sha256=sd2cUdRJG-UCdQr9ruZvEYpNGNF1t2W2fuxwwVlQD9E,4164
 nshtrainer/trainer/checkpoint_connector.py,sha256=F2tkHogbMAa5U7335sm77sZBkjEDa5v46XbJCH9Mg6c,2167
 nshtrainer/trainer/signal_connector.py,sha256=llwc8pdKAWxREFpjdi14Bpy8rGVMEJsmJx_s2p4gI8E,10689
-nshtrainer/trainer/trainer.py,sha256=IHEtuDVVBradVQOKSP9zYAalkn2sguXUZixzvS8P4UY,17097
+nshtrainer/trainer/trainer.py,sha256=M97phnALfG18VxkMLoDr5AKFf4UaPBdc6S2BghdBtas,17103
 nshtrainer/util/_environment_info.py,sha256=yPtAbgjCY4tkvh5wp9sjNsF0Z45TYwzEAM_N2_b5BbY,23123
 nshtrainer/util/_useful_types.py,sha256=dwZokFkIe7M5i2GR3nQ9A1lhGw06DMAFfH5atyquqSA,8000
 nshtrainer/util/environment.py,sha256=AeW_kLl-N70wmb6L_JLz1wRj0kA70xs6RCmc9iUqczE,4159
@@ -82,6 +87,6 @@ nshtrainer/util/seed.py,sha256=Or2wMPsnQxfnZ2xfBiyMcHFIUt3tGTNeMMyOEanCkqs,280
 nshtrainer/util/slurm.py,sha256=rofIU26z3SdL79SF45tNez6juou1cyDLz07oXEZb9Hg,1566
 nshtrainer/util/typed.py,sha256=NGuDkDzFlc1fAoaXjOFZVbmj0mRFjsQi1E_hPa7Bn5U,128
 nshtrainer/util/typing_utils.py,sha256=8ptjSSLZxlmy4FY6lzzkoGoF5fGNClo8-B_c0XHQaNU,385
-nshtrainer-0.13.1.dist-info/METADATA,sha256=BKPztr_MCloMTaIsrPNhk8PZNhKgTIED7Lep-8JAi1s,860
-nshtrainer-0.13.1.dist-info/WHEEL,sha256=sP946D7jFCHeNz5Iq4fL4Lu-PrWrFsgfLXbbkciIZwg,88
-nshtrainer-0.13.1.dist-info/RECORD,,
+nshtrainer-0.14.1.dist-info/METADATA,sha256=V7kssNBEV-JqILCZbRujHMs4ORomy_hWU4vRK0Je2GU,860
+nshtrainer-0.14.1.dist-info/WHEEL,sha256=sP946D7jFCHeNz5Iq4fL4Lu-PrWrFsgfLXbbkciIZwg,88
+nshtrainer-0.14.1.dist-info/RECORD,,