nshtrainer 1.0.0b45__py3-none-any.whl → 1.0.0b47__py3-none-any.whl

nshtrainer/_callback.py

@@ -4,46 +4,46 @@ from pathlib import Path
 from typing import TYPE_CHECKING, Any
 
 import torch
+from lightning.pytorch import LightningModule
 from lightning.pytorch.callbacks import Callback as _LightningCallback
 from lightning.pytorch.utilities.types import STEP_OUTPUT
 from torch.optim import Optimizer
 
 if TYPE_CHECKING:
-    from .model import LightningModuleBase
     from .trainer import Trainer
 
 
 class NTCallbackBase(_LightningCallback):
     def setup(  # pyright: ignore[reportIncompatibleMethodOverride]
-        self, trainer: Trainer, pl_module: LightningModuleBase, stage: str
+        self, trainer: Trainer, pl_module: LightningModule, stage: str
     ) -> None:
         """Called when fit, validate, test, predict, or tune begins."""
 
     def teardown(  # pyright: ignore[reportIncompatibleMethodOverride]
-        self, trainer: Trainer, pl_module: LightningModuleBase, stage: str
+        self, trainer: Trainer, pl_module: LightningModule, stage: str
     ) -> None:
         """Called when fit, validate, test, predict, or tune ends."""
 
-    def on_fit_start(self, trainer: Trainer, pl_module: LightningModuleBase) -> None:  # pyright: ignore[reportIncompatibleMethodOverride]
+    def on_fit_start(self, trainer: Trainer, pl_module: LightningModule) -> None:  # pyright: ignore[reportIncompatibleMethodOverride]
         """Called when fit begins."""
 
-    def on_fit_end(self, trainer: Trainer, pl_module: LightningModuleBase) -> None:  # pyright: ignore[reportIncompatibleMethodOverride]
+    def on_fit_end(self, trainer: Trainer, pl_module: LightningModule) -> None:  # pyright: ignore[reportIncompatibleMethodOverride]
         """Called when fit ends."""
 
     def on_sanity_check_start(  # pyright: ignore[reportIncompatibleMethodOverride]
-        self, trainer: Trainer, pl_module: LightningModuleBase
+        self, trainer: Trainer, pl_module: LightningModule
     ) -> None:
         """Called when the validation sanity check starts."""
 
     def on_sanity_check_end(  # pyright: ignore[reportIncompatibleMethodOverride]
-        self, trainer: Trainer, pl_module: LightningModuleBase
+        self, trainer: Trainer, pl_module: LightningModule
     ) -> None:
         """Called when the validation sanity check ends."""
 
     def on_train_batch_start(  # pyright: ignore[reportIncompatibleMethodOverride]
         self,
         trainer: Trainer,
-        pl_module: LightningModuleBase,
+        pl_module: LightningModule,
         batch: Any,
         batch_idx: int,
     ) -> None:
@@ -52,7 +52,7 @@ class NTCallbackBase(_LightningCallback):
     def on_train_batch_end(  # pyright: ignore[reportIncompatibleMethodOverride]
         self,
         trainer: Trainer,
-        pl_module: LightningModuleBase,
+        pl_module: LightningModule,
         outputs: STEP_OUTPUT,
         batch: Any,
         batch_idx: int,
@@ -66,12 +66,12 @@ class NTCallbackBase(_LightningCallback):
         """
 
     def on_train_epoch_start(  # pyright: ignore[reportIncompatibleMethodOverride]
-        self, trainer: Trainer, pl_module: LightningModuleBase
+        self, trainer: Trainer, pl_module: LightningModule
     ) -> None:
         """Called when the train epoch begins."""
 
     def on_train_epoch_end(  # pyright: ignore[reportIncompatibleMethodOverride]
-        self, trainer: Trainer, pl_module: LightningModuleBase
+        self, trainer: Trainer, pl_module: LightningModule
     ) -> None:
         """Called when the train epoch ends.
 
@@ -102,39 +102,39 @@ class NTCallbackBase(_LightningCallback):
         """
 
     def on_validation_epoch_start(  # pyright: ignore[reportIncompatibleMethodOverride]
-        self, trainer: Trainer, pl_module: LightningModuleBase
+        self, trainer: Trainer, pl_module: LightningModule
     ) -> None:
         """Called when the val epoch begins."""
 
     def on_validation_epoch_end(  # pyright: ignore[reportIncompatibleMethodOverride]
-        self, trainer: Trainer, pl_module: LightningModuleBase
+        self, trainer: Trainer, pl_module: LightningModule
     ) -> None:
         """Called when the val epoch ends."""
 
     def on_test_epoch_start(  # pyright: ignore[reportIncompatibleMethodOverride]
-        self, trainer: Trainer, pl_module: LightningModuleBase
+        self, trainer: Trainer, pl_module: LightningModule
     ) -> None:
         """Called when the test epoch begins."""
 
     def on_test_epoch_end(  # pyright: ignore[reportIncompatibleMethodOverride]
-        self, trainer: Trainer, pl_module: LightningModuleBase
+        self, trainer: Trainer, pl_module: LightningModule
     ) -> None:
         """Called when the test epoch ends."""
 
     def on_predict_epoch_start(  # pyright: ignore[reportIncompatibleMethodOverride]
-        self, trainer: Trainer, pl_module: LightningModuleBase
+        self, trainer: Trainer, pl_module: LightningModule
     ) -> None:
         """Called when the predict epoch begins."""
 
     def on_predict_epoch_end(  # pyright: ignore[reportIncompatibleMethodOverride]
-        self, trainer: Trainer, pl_module: LightningModuleBase
+        self, trainer: Trainer, pl_module: LightningModule
     ) -> None:
         """Called when the predict epoch ends."""
 
     def on_validation_batch_start(  # pyright: ignore[reportIncompatibleMethodOverride]
         self,
         trainer: Trainer,
-        pl_module: LightningModuleBase,
+        pl_module: LightningModule,
         batch: Any,
         batch_idx: int,
         dataloader_idx: int = 0,
@@ -144,7 +144,7 @@ class NTCallbackBase(_LightningCallback):
     def on_validation_batch_end(  # pyright: ignore[reportIncompatibleMethodOverride]
         self,
         trainer: Trainer,
-        pl_module: LightningModuleBase,
+        pl_module: LightningModule,
         outputs: STEP_OUTPUT,
         batch: Any,
         batch_idx: int,
@@ -155,7 +155,7 @@ class NTCallbackBase(_LightningCallback):
     def on_test_batch_start(  # pyright: ignore[reportIncompatibleMethodOverride]
         self,
         trainer: Trainer,
-        pl_module: LightningModuleBase,
+        pl_module: LightningModule,
         batch: Any,
         batch_idx: int,
         dataloader_idx: int = 0,
@@ -165,7 +165,7 @@ class NTCallbackBase(_LightningCallback):
     def on_test_batch_end(  # pyright: ignore[reportIncompatibleMethodOverride]
         self,
         trainer: Trainer,
-        pl_module: LightningModuleBase,
+        pl_module: LightningModule,
         outputs: STEP_OUTPUT,
         batch: Any,
         batch_idx: int,
@@ -176,7 +176,7 @@ class NTCallbackBase(_LightningCallback):
     def on_predict_batch_start(  # pyright: ignore[reportIncompatibleMethodOverride]
         self,
         trainer: Trainer,
-        pl_module: LightningModuleBase,
+        pl_module: LightningModule,
         batch: Any,
         batch_idx: int,
         dataloader_idx: int = 0,
@@ -186,7 +186,7 @@ class NTCallbackBase(_LightningCallback):
     def on_predict_batch_end(  # pyright: ignore[reportIncompatibleMethodOverride]
         self,
         trainer: Trainer,
-        pl_module: LightningModuleBase,
+        pl_module: LightningModule,
         outputs: Any,
         batch: Any,
         batch_idx: int,
@@ -194,40 +194,40 @@ class NTCallbackBase(_LightningCallback):
     ) -> None:
         """Called when the predict batch ends."""
 
-    def on_train_start(self, trainer: Trainer, pl_module: LightningModuleBase) -> None:  # pyright: ignore[reportIncompatibleMethodOverride]
+    def on_train_start(self, trainer: Trainer, pl_module: LightningModule) -> None:  # pyright: ignore[reportIncompatibleMethodOverride]
         """Called when the train begins."""
 
-    def on_train_end(self, trainer: Trainer, pl_module: LightningModuleBase) -> None:  # pyright: ignore[reportIncompatibleMethodOverride]
+    def on_train_end(self, trainer: Trainer, pl_module: LightningModule) -> None:  # pyright: ignore[reportIncompatibleMethodOverride]
         """Called when the train ends."""
 
     def on_validation_start(  # pyright: ignore[reportIncompatibleMethodOverride]
-        self, trainer: Trainer, pl_module: LightningModuleBase
+        self, trainer: Trainer, pl_module: LightningModule
     ) -> None:
         """Called when the validation loop begins."""
 
     def on_validation_end(  # pyright: ignore[reportIncompatibleMethodOverride]
-        self, trainer: Trainer, pl_module: LightningModuleBase
+        self, trainer: Trainer, pl_module: LightningModule
     ) -> None:
         """Called when the validation loop ends."""
 
-    def on_test_start(self, trainer: Trainer, pl_module: LightningModuleBase) -> None:  # pyright: ignore[reportIncompatibleMethodOverride]
+    def on_test_start(self, trainer: Trainer, pl_module: LightningModule) -> None:  # pyright: ignore[reportIncompatibleMethodOverride]
         """Called when the test begins."""
 
-    def on_test_end(self, trainer: Trainer, pl_module: LightningModuleBase) -> None:  # pyright: ignore[reportIncompatibleMethodOverride]
+    def on_test_end(self, trainer: Trainer, pl_module: LightningModule) -> None:  # pyright: ignore[reportIncompatibleMethodOverride]
         """Called when the test ends."""
 
     def on_predict_start(  # pyright: ignore[reportIncompatibleMethodOverride]
-        self, trainer: Trainer, pl_module: LightningModuleBase
+        self, trainer: Trainer, pl_module: LightningModule
     ) -> None:
         """Called when the predict begins."""
 
-    def on_predict_end(self, trainer: Trainer, pl_module: LightningModuleBase) -> None:  # pyright: ignore[reportIncompatibleMethodOverride]
+    def on_predict_end(self, trainer: Trainer, pl_module: LightningModule) -> None:  # pyright: ignore[reportIncompatibleMethodOverride]
         """Called when predict ends."""
 
     def on_exception(  # pyright: ignore[reportIncompatibleMethodOverride]
         self,
         trainer: Trainer,
-        pl_module: LightningModuleBase,
+        pl_module: LightningModule,
         exception: BaseException,
     ) -> None:
         """Called when any trainer execution is interrupted by an exception."""
@@ -253,7 +253,7 @@ class NTCallbackBase(_LightningCallback):
     def on_save_checkpoint(  # pyright: ignore[reportIncompatibleMethodOverride]
         self,
         trainer: Trainer,
-        pl_module: LightningModuleBase,
+        pl_module: LightningModule,
         checkpoint: dict[str, Any],
     ) -> None:
         r"""Called when saving a checkpoint to give you a chance to store anything else you might want to save.
@@ -268,7 +268,7 @@ class NTCallbackBase(_LightningCallback):
     def on_load_checkpoint(  # pyright: ignore[reportIncompatibleMethodOverride]
         self,
         trainer: Trainer,
-        pl_module: LightningModuleBase,
+        pl_module: LightningModule,
         checkpoint: dict[str, Any],
     ) -> None:
         r"""Called when loading a model checkpoint, use to reload state.
@@ -281,19 +281,19 @@ class NTCallbackBase(_LightningCallback):
         """
 
     def on_before_backward(  # pyright: ignore[reportIncompatibleMethodOverride]
-        self, trainer: Trainer, pl_module: LightningModuleBase, loss: torch.Tensor
+        self, trainer: Trainer, pl_module: LightningModule, loss: torch.Tensor
     ) -> None:
         """Called before ``loss.backward()``."""
 
     def on_after_backward(  # pyright: ignore[reportIncompatibleMethodOverride]
-        self, trainer: Trainer, pl_module: LightningModuleBase
+        self, trainer: Trainer, pl_module: LightningModule
     ) -> None:
         """Called after ``loss.backward()`` and before optimizers are stepped."""
 
     def on_before_optimizer_step(  # pyright: ignore[reportIncompatibleMethodOverride]
         self,
         trainer: Trainer,
-        pl_module: LightningModuleBase,
+        pl_module: LightningModule,
         optimizer: Optimizer,
     ) -> None:
         """Called before ``optimizer.step()``."""
@@ -301,7 +301,7 @@ class NTCallbackBase(_LightningCallback):
     def on_before_zero_grad(  # pyright: ignore[reportIncompatibleMethodOverride]
         self,
         trainer: Trainer,
-        pl_module: LightningModuleBase,
+        pl_module: LightningModule,
         optimizer: Optimizer,
     ) -> None:
         """Called before ``optimizer.zero_grad()``."""
@@ -310,15 +310,15 @@ class NTCallbackBase(_LightningCallback):
         self,
         ckpt_path: Path,
         metadata_path: Path | None,
-        trainer: "Trainer",
-        pl_module: "LightningModuleBase",
+        trainer: Trainer,
+        pl_module: LightningModule,
     ) -> None:
         """Called after a checkpoint is saved."""
         pass
 
 
 def _call_on_checkpoint_saved(
-    trainer: "Trainer",
+    trainer: Trainer,
     ckpt_path: str | Path,
     metadata_path: str | Path | None,
 ):
@@ -333,5 +333,5 @@ def _call_on_checkpoint_saved(
             ckpt_path,
             metadata_path,
             trainer,
-            trainer._base_module,
+            trainer.lightning_module,
         )

nshtrainer/_checkpoint/metadata.py

@@ -16,6 +16,18 @@ from ..util.path import compute_file_checksum, try_symlink_or_copy
 if TYPE_CHECKING:
     from ..trainer.trainer import Trainer
 
+try:
+    from pydantic import BaseModel
+
+    _HAS_PYDANTIC = True
+except ImportError:
+    if not TYPE_CHECKING:
+        BaseModel = object
+    else:
+        from pydantic import BaseModel
+    _HAS_PYDANTIC = False
+
+
 log = logging.getLogger(__name__)
 
 
@@ -27,10 +39,10 @@ def _full_hparams_dict(trainer: Trainer):
     hparams["trainer"] = trainer.hparams.model_dump(mode="json")
 
     if trainer.lightning_module is not None:
-        from ..model import LightningModuleBase
-
-        if isinstance(trainer.lightning_module, LightningModuleBase):
-            hparams["model"] = trainer.lightning_module.hparams.model_dump(mode="json")
+        model_hparams = trainer.lightning_module.hparams
+        if _HAS_PYDANTIC and isinstance(model_hparams, BaseModel):
+            model_hparams = model_hparams.model_dump(mode="json")
+        hparams["model"] = dict(model_hparams)
 
     return hparams
 
@@ -139,15 +151,30 @@ def remove_checkpoint_metadata(checkpoint_path: Path):
         log.debug(f"Removed {path}")
 
 
+def remove_checkpoint_metadata_link(ckpt_link_path: Path):
+    path = _metadata_path(ckpt_link_path)
+    # If the metadata does not exist, we can safely ignore this
+    if not path.exists(follow_symlinks=False):
+        # This is EXTREMELY important here
+        # Otherwise, we've already deleted the file that the symlink
+        # used to point to, so this always returns False
+        log.debug(f"Metadata file does not exist: {path}")
+        return
+
+    # If the metadata exists, we can remove it
+    try:
+        path.unlink(missing_ok=True)
+    except Exception:
+        log.warning(f"Failed to remove {path}", exc_info=True)
+    else:
+        log.debug(f"Removed {path}")
+
+
 def link_checkpoint_metadata(checkpoint_path: Path, linked_checkpoint_path: Path):
     # First, remove any existing metadata files
-    remove_checkpoint_metadata(linked_checkpoint_path)
+    remove_checkpoint_metadata_link(linked_checkpoint_path)
 
     # Link the metadata files to the new checkpoint
     path = _metadata_path(checkpoint_path)
     linked_path = _metadata_path(linked_checkpoint_path)
-
-    if not path.exists():
-        raise FileNotFoundError(f"Checkpoint path does not exist: {checkpoint_path}")
-
     try_symlink_or_copy(path, linked_path)

nshtrainer/_checkpoint/saver.py

@@ -8,7 +8,11 @@ from pathlib import Path
 from lightning.pytorch import Trainer
 
 from ..util.path import try_symlink_or_copy
-from .metadata import link_checkpoint_metadata, remove_checkpoint_metadata
+from .metadata import (
+    link_checkpoint_metadata,
+    remove_checkpoint_metadata,
+    remove_checkpoint_metadata_link,
+)
 
 log = logging.getLogger(__name__)
 
@@ -39,7 +43,7 @@ def link_checkpoint(
             log.debug(f"Removed {linkpath=}")
 
         if metadata:
-            remove_checkpoint_metadata(linkpath)
+            remove_checkpoint_metadata_link(linkpath)
 
     try_symlink_or_copy(filepath, linkpath)
     if metadata:

nshtrainer/callbacks/checkpoint/_base.py

@@ -160,7 +160,7 @@ class CheckpointBase(Checkpoint, ABC, Generic[TConfig]):
         filepath = self.resolve_checkpoint_path(self.current_metrics(trainer))
         trainer.save_checkpoint(filepath, self.config.save_weights_only)
 
-        if trainer.is_global_zero:
+        if trainer.hparams.save_checkpoint_metadata and trainer.is_global_zero:
             # Remove old checkpoints
             self.remove_old_checkpoints(trainer)
 

nshtrainer/lr_scheduler/base.py

@@ -3,9 +3,10 @@ from __future__ import annotations
 import math
 from abc import ABC, abstractmethod
 from collections.abc import Mapping
-from typing import TYPE_CHECKING, Literal
+from typing import Literal
 
 import nshconfig as C
+from lightning.pytorch import LightningModule
 from lightning.pytorch.utilities.types import (
     LRSchedulerConfigType,
     LRSchedulerTypeUnion,
@@ -13,9 +14,6 @@ from lightning.pytorch.utilities.types import (
 from torch.optim import Optimizer
 from typing_extensions import Never, NotRequired, TypedDict
 
-if TYPE_CHECKING:
-    from ..model.base import LightningModuleBase
-
 
 class LRSchedulerMetadata(TypedDict):
     interval: Literal["epoch", "step"]
@@ -44,13 +42,13 @@ class LRSchedulerConfigBase(C.Config, ABC):
 
     @abstractmethod
     def create_scheduler_impl(
-        self, optimizer: Optimizer, lightning_module: LightningModuleBase
+        self, optimizer: Optimizer, lightning_module: LightningModule
     ) -> LRSchedulerTypeUnion | LRSchedulerConfigType: ...
 
     def create_scheduler(
         self,
         optimizer: Optimizer,
-        lightning_module: LightningModuleBase,
+        lightning_module: LightningModule,
         lr: Never
         | None = None,  # Backward compatibility, should be removed in the future
     ) -> LRSchedulerConfigType:
@@ -87,7 +85,7 @@ class LRSchedulerConfigBase(C.Config, ABC):
 
         return scheduler
 
-    def compute_num_steps_per_epoch(self, lightning_module: LightningModuleBase) -> int:
+    def compute_num_steps_per_epoch(self, lightning_module: LightningModule) -> int:
         trainer = lightning_module.trainer
         # Use the Lightning trainer to convert the epoch-based values to step-based values
         _ = trainer.estimated_stepping_batches

nshtrainer/model/base.py

@@ -62,73 +62,6 @@ class LightningModuleBase(
     ABC,
     Generic[THparams],
 ):
-    # region Debug
-    @property
-    def debug(self) -> bool:
-        if torch.jit.is_scripting():
-            return False
-
-        if (trainer := self._trainer) is None:
-            return False
-
-        from ..trainer import Trainer
-
-        if not isinstance(trainer, Trainer):
-            return False
-
-        return trainer.debug
-
-    @debug.setter
-    def debug(self, value: bool):
-        if torch.jit.is_scripting():
-            return
-
-        if (trainer := self._trainer) is None:
-            return
-
-        from ..trainer import Trainer
-
-        if not isinstance(trainer, Trainer):
-            return
-
-        trainer.debug = value
-
-    @torch.jit.unused
-    def breakpoint(self, rank_zero_only: bool = True):
-        if (
-            not rank_zero_only
-            or not torch.distributed.is_initialized()
-            or torch.distributed.get_rank() == 0
-        ):
-            breakpoint()
-
-        if rank_zero_only and torch.distributed.is_initialized():
-            _ = torch.distributed.barrier()
-
-    @torch.jit.unused
-    def ensure_finite(
-        self,
-        tensor: torch.Tensor,
-        name: str | None = None,
-        throw: bool = False,
-    ):
-        name_parts: list[str] = ["Tensor"]
-        if name is not None:
-            name_parts.append(name)
-        name = " ".join(name_parts)
-
-        not_finite = ~torch.isfinite(tensor)
-        if not_finite.any():
-            msg = f"{name} has {not_finite.sum().item()}/{not_finite.numel()} non-finite values."
-            if throw:
-                raise RuntimeError(msg)
-            else:
-                log.warning(msg)
-            return False
-        return True
-
-    # endregion
-
     # region Profiler
     @property
     def profiler(self) -> Profiler:

nshtrainer/trainer/trainer.py

@@ -418,20 +418,6 @@ class Trainer(LightningTrainer):
 
         return tracker.time_elapsed(stage)
 
-    @property
-    def _base_module(self):
-        if self.lightning_module is None:
-            raise ValueError("LightningModule is not set.")
-
-        from ..model.base import LightningModuleBase
-
-        if not isinstance(self.lightning_module, LightningModuleBase):
-            raise ValueError(
-                f"LightningModule is not an instance of {LightningModuleBase}."
-            )
-
-        return self.lightning_module
-
     @override
     def _run(
         self, model: LightningModule, ckpt_path: str | Path | None = None

nshtrainer/util/path.py

@@ -120,7 +120,8 @@ def try_symlink_or_copy(
                 shutil.copy(file_path, link_path)
         else:
             link_path.symlink_to(
-                symlink_target, target_is_directory=target_is_directory
+                symlink_target,
+                target_is_directory=target_is_directory,
             )
     except Exception:
         log.warning(

{nshtrainer-1.0.0b45.dist-info → nshtrainer-1.0.0b47.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.3
 Name: nshtrainer
-Version: 1.0.0b45
+Version: 1.0.0b47
 Summary: 
 Author: Nima Shoghi
 Author-email: nimashoghi@gmail.com

{nshtrainer-1.0.0b45.dist-info → nshtrainer-1.0.0b47.dist-info}/RECORD

@@ -1,8 +1,8 @@
 nshtrainer/.nshconfig.generated.json,sha256=yZd6cn1RhvNNJUgiUTRYut8ofZYvbulnpPG-rZIRhi4,106
 nshtrainer/__init__.py,sha256=g_moPnfQxSxFZX5NB9ILQQOJrt4RTRuiFt9N0STIpxM,874
-nshtrainer/_callback.py,sha256=tXQCDzS6CvMTuTY5lQSH5qZs1pXUi-gt9bQdpXMVdEs,12715
-nshtrainer/_checkpoint/metadata.py,sha256=LQZ8g50rKxQQx-FqiW3n8EWmal9qSWRouOpIIn6NJJY,4758
-nshtrainer/_checkpoint/saver.py,sha256=rWl4d2lCTMU4_wt8yZFL2pFQaP9hj5sPgqHMPQ4zuyI,1584
+nshtrainer/_callback.py,sha256=ZDppiJ4d65tRXTEWYPZLH_F1xFizdz1pkWJe_sQ5uII,12564
+nshtrainer/_checkpoint/metadata.py,sha256=ojSEmq0udFwdzIC5vkbF0yEdhMaJ2iBrZCSFNDkeeGY,5578
+nshtrainer/_checkpoint/saver.py,sha256=65UDrz3KuhkgVfco-RkWuoa1wzTZoXxunlC769yJaMc,1639
 nshtrainer/_directory.py,sha256=TJR9ccyuzRlAVfVjGyeQ3E2AFAcz-XbBCxWfiXo2SlY,3191
 nshtrainer/_experimental/__init__.py,sha256=U4S_2y3zgLZVfMenHRaJFBW8yqh2mUBuI291LGQVOJ8,35
 nshtrainer/_hf_hub.py,sha256=4OsCbIITnZk_YLyoMrVyZ0SIN04FBxlC0ig2Et8UAdo,14287
@@ -10,7 +10,7 @@ nshtrainer/callbacks/__init__.py,sha256=w80d6PGNu3wjUj9NiRGMqCX9NnXD5ZlvbY-DIK4z
 nshtrainer/callbacks/actsave.py,sha256=NSXIIu62MNYe5gz479SMW33bdoKYoYtWtd_iTWFpKpc,3881
 nshtrainer/callbacks/base.py,sha256=Alaou1IHAIlMEM7g58d_02ozY2xWlshBN7fsw5Ee21s,3683
 nshtrainer/callbacks/checkpoint/__init__.py,sha256=l8tkHc83_mLiU0-wT09SWdRzwpm2ulbkLzcuCmuTwzE,620
-nshtrainer/callbacks/checkpoint/_base.py,sha256=wCJBRI0pQYZc3GBu0b-aUBlBDhd39AdL82VvFgKmv3k,6300
+nshtrainer/callbacks/checkpoint/_base.py,sha256=f7lpk8W4xqxk3PolBEU3AWt9VTIpoLW7wMUhC5DNm3c,6345
 nshtrainer/callbacks/checkpoint/best_checkpoint.py,sha256=2CQuhPJ3Fi7lDw7z-J8kXXXuDU8-4HcU48oZxR49apk,2667
 nshtrainer/callbacks/checkpoint/last_checkpoint.py,sha256=vn-as3ex7kaTRcKsIurVtM6kUSHYNwHJeYG82j2dMcc,3554
 nshtrainer/callbacks/checkpoint/on_exception_checkpoint.py,sha256=nljzETqkHwA-4g8mxaeFK5HxA8My0dlIPzIUscSMWyk,3525
@@ -109,13 +109,13 @@ nshtrainer/loggers/csv.py,sha256=xJ8mSmw4vJwinIfqhF6t2HWmh_1dXEYyLfGuXwL7WHo,116
 nshtrainer/loggers/tensorboard.py,sha256=E7iO_fDt9bfH02hBL430bXPLljOo5iGgq2QyPqmx2gQ,2324
 nshtrainer/loggers/wandb.py,sha256=KZXAUWrrmdX_L8rqej77oUHaM0JxZRM8y9z6JP9PISw,6856
 nshtrainer/lr_scheduler/__init__.py,sha256=daMMK3erUcNXGGd_nZB8AWu3ZTYqfS1RSWeK4FV2udw,851
-nshtrainer/lr_scheduler/base.py,sha256=062fGcH5sYeEKwoY55RydCTvfPwTnyZHCi049a3nMbM,3805
+nshtrainer/lr_scheduler/base.py,sha256=LE53JRBTuAlA1fqbMgCZ7m39D1z0rGj2TizhJ62CPvE,3756
 nshtrainer/lr_scheduler/linear_warmup_cosine.py,sha256=MsoXgCcWTKsrkNZiGnKS6yC-slRuleuwFxeM_lmG_pQ,5560
 nshtrainer/lr_scheduler/reduce_lr_on_plateau.py,sha256=v9T0GpvOoHV30atFB0MwExHgHcTpMCYxbMRoPjPBjt8,2938
 nshtrainer/metrics/__init__.py,sha256=Nqkn_jsDf3n5WtfMcnaaEftYjIIT2b-S7rmsB1MOMkU,86
 nshtrainer/metrics/_config.py,sha256=XIRokFM8PHrhBa3w2R6BM6a4es3ncsoBqE_LqXQFsFE,1223
 nshtrainer/model/__init__.py,sha256=3G-bwPPSRStWdsdwG9-rn0bXcRpEiP1BiQpF_qavtls,97
-nshtrainer/model/base.py,sha256=JL3AmH17GQjQIoMrZl3O0vUI7dj5ZsO5iEJgoLPyzHw,10356
+nshtrainer/model/base.py,sha256=q1IMVG3lHvI84x-8hXmiLNJN_NplY_q9W5u6D2rrmVY,8684
 nshtrainer/model/mixins/callback.py,sha256=0LPgve4VszHbLipid4mpI1qnnmdGS2spivs0dXLvqHw,3154
 nshtrainer/model/mixins/debug.py,sha256=1LX9KzeFX9JDPs_a6YCdYDZXLhEk_5rBO2aCqlfBy7w,2087
 nshtrainer/model/mixins/logger.py,sha256=27H99FuLaxc6_dDLG2pid4E_5E0-eLGnc2Ifpt0HYIM,6066
@@ -142,18 +142,18 @@ nshtrainer/trainer/plugin/layer_sync.py,sha256=h-ydZwXepnsw5-paLgiDatqPyQ_8C0QEv
 nshtrainer/trainer/plugin/precision.py,sha256=I0QsB1bVxmsFmBOkgrAfGONsuYae_lD9Bz0PfJEQvH4,5598
 nshtrainer/trainer/signal_connector.py,sha256=GhfGcSzfaTNhnj2QFkBDq5aT7FqbLMA7eC8SYQs8_8w,10828
 nshtrainer/trainer/strategy.py,sha256=VPTn5z3zvXTydY8IJchjhjcOfpvtoejnvUkq5E4WTus,1368
-nshtrainer/trainer/trainer.py,sha256=8wMe0qArbDfStS4UdmuKSC2aiAImR3mhj14_kCJiNSM,20797
+nshtrainer/trainer/trainer.py,sha256=QEK-0bcw1y5Cconi99PYFXr0MElUGgGYMZ_SlcJUQ1k,20364
 nshtrainer/util/_environment_info.py,sha256=MT8mBe6ZolRfKiwU-les1P-lPNPqXpHQcfADrh_A3uY,24629
 nshtrainer/util/bf16.py,sha256=9QhHZCkYSfYpIcxwAMoXyuh2yTSHBzT-EdLQB297jEs,762
 nshtrainer/util/config/__init__.py,sha256=Z39JJufSb61Lhn2GfVcv3eFW_eorOrN9-9llDWlnZZM,272
 nshtrainer/util/config/dtype.py,sha256=Fn_MhhQoHPyFAnFPSwvcvLiGR3yWFIszMba02CJiC4g,2213
 nshtrainer/util/config/duration.py,sha256=mM-UfU_HvhXwW33TYEDg0x58n80tnle2e6VaWtxZTjk,764
 nshtrainer/util/environment.py,sha256=s-B5nY0cKYXdFMdNYumvC_xxacMATiI4DvV2gUDu20k,4195
-nshtrainer/util/path.py,sha256=9fIjE3S78pPL6wjAgEJUYfIJQAPdKOQqIYvTS9lWTUk,3959
+nshtrainer/util/path.py,sha256=wADa_qFiekiX0PENf9jI-9_abICB0vJFXav8pNLR7lw,3976
 nshtrainer/util/seed.py,sha256=diMV8iwBKN7Xxt5pELmui-gyqyT80_CZzomrWhNss0k,316
 nshtrainer/util/slurm.py,sha256=HflkP5iI_r4UHMyPjw9R4dD5AHsJUpcfJw5PLvGYBRM,1603
 nshtrainer/util/typed.py,sha256=Xt5fUU6zwLKSTLUdenovnKK0N8qUq89Kddz2_XeykVQ,164
 nshtrainer/util/typing_utils.py,sha256=MjY-CUX9R5Tzat-BlFnQjwl1PQ_W2yZQoXhkYHlJ_VA,442
-nshtrainer-1.0.0b45.dist-info/METADATA,sha256=_RPpe6F7DXpsQSmBF1GTc-E5VUfaC69fIYfoFhsip2s,988
-nshtrainer-1.0.0b45.dist-info/WHEEL,sha256=XbeZDeTWKc1w7CSIyre5aMDU_-PohRwTQceYnisIYYY,88
-nshtrainer-1.0.0b45.dist-info/RECORD,,
+nshtrainer-1.0.0b47.dist-info/METADATA,sha256=E7d5EfVnqLgmFPuh_D_VWERKQrA5tjePktx1vujkSs8,988
+nshtrainer-1.0.0b47.dist-info/WHEEL,sha256=XbeZDeTWKc1w7CSIyre5aMDU_-PohRwTQceYnisIYYY,88
+nshtrainer-1.0.0b47.dist-info/RECORD,,