nshtrainer 0.41.1__tar.gz → 0.43.0__tar.gz

{nshtrainer-0.41.1 → nshtrainer-0.43.0}/PKG-INFO

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

{nshtrainer-0.41.1 → nshtrainer-0.43.0}/pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "nshtrainer"
-version = "0.41.1"
+version = "0.43.0"
 description = ""
 authors = ["Nima Shoghi <nimashoghi@gmail.com>"]
 readme = "README.md"
@@ -25,10 +25,10 @@ tensorboard = { version = "*", optional = true }
 huggingface-hub = { version = "*", optional = true }
 
 [tool.poetry.group.dev.dependencies]
-pyright = "^1.1.372"
-ruff = "^0.5.4"
-ipykernel = "^6.29.5"
-ipywidgets = "^8.1.3"
+pyright = "*"
+ruff = "*"
+ipykernel = "*"
+ipywidgets = "*"
 
 [build-system]
 requires = ["poetry-core"]
@@ -43,7 +43,11 @@ strictSetInference = true
 reportPrivateImportUsage = false
 
 [tool.ruff.lint]
+select = ["FA102", "FA100"]
 ignore = ["F722", "F821", "E731", "E741"]
 
+[tool.ruff.lint.isort]
+required-imports = ["from __future__ import annotations"]
+
 [tool.poetry.extras]
 extra = ["wrapt", "GitPython", "wandb", "tensorboard", "huggingface-hub"]

{nshtrainer-0.41.1 → nshtrainer-0.43.0}/src/nshtrainer/__init__.py

@@ -1,3 +1,5 @@
+from __future__ import annotations
+
 from . import _experimental as _experimental
 from . import callbacks as callbacks
 from . import config as config

{nshtrainer-0.41.1 → nshtrainer-0.43.0}/src/nshtrainer/_callback.py

@@ -1,3 +1,5 @@
+from __future__ import annotations
+
 from pathlib import Path
 from typing import TYPE_CHECKING
 

{nshtrainer-0.41.1 → nshtrainer-0.43.0}/src/nshtrainer/_checkpoint/loader.py

@@ -1,3 +1,5 @@
+from __future__ import annotations
+
 import logging
 from collections.abc import Iterable, Sequence
 from dataclasses import dataclass

{nshtrainer-0.41.1 → nshtrainer-0.43.0}/src/nshtrainer/_checkpoint/metadata.py

@@ -1,3 +1,5 @@
+from __future__ import annotations
+
 import copy
 import datetime
 import logging

{nshtrainer-0.41.1 → nshtrainer-0.43.0}/src/nshtrainer/_checkpoint/saver.py

@@ -1,3 +1,5 @@
+from __future__ import annotations
+
 import logging
 import os
 import shutil

{nshtrainer-0.41.1 → nshtrainer-0.43.0}/src/nshtrainer/_directory.py

@@ -1,9 +1,11 @@
+from __future__ import annotations
+
 import logging
 from pathlib import Path
 
 import nshconfig as C
 
-from .callbacks.directory_setup import DirectorySetupConfig
+from .callbacks.directory_setup import DirectorySetupCallbackConfig
 from .loggers import LoggerConfig
 
 log = logging.getLogger(__name__)
@@ -32,7 +34,7 @@ class DirectoryConfig(C.Config):
     profile: Path | None = None
     """Directory to save profiling information to. If None, will use nshtrainer/{id}/profile/."""
 
-    setup_callback: DirectorySetupConfig = DirectorySetupConfig()
+    setup_callback: DirectorySetupCallbackConfig = DirectorySetupCallbackConfig()
     """Configuration for the directory setup PyTorch Lightning callback."""
 
     def resolve_run_root_directory(self, run_id: str) -> Path:

{nshtrainer-0.41.1 → nshtrainer-0.43.0}/src/nshtrainer/_experimental/__init__.py

@@ -1 +1,3 @@
+from __future__ import annotations
+
 from lightning.fabric.utilities.throughput import measure_flops as measure_flops

{nshtrainer-0.41.1 → nshtrainer-0.43.0}/src/nshtrainer/_hf_hub.py

@@ -1,3 +1,5 @@
+from __future__ import annotations
+
 import contextlib
 import logging
 import os

nshtrainer-0.43.0/src/nshtrainer/callbacks/__init__.py

@@ -0,0 +1,81 @@
+from __future__ import annotations
+
+from typing import Annotated
+
+import nshconfig as C
+
+from . import checkpoint as checkpoint
+from .base import CallbackConfigBase as CallbackConfigBase
+from .checkpoint import BestCheckpoint as BestCheckpoint
+from .checkpoint import BestCheckpointCallbackConfig as BestCheckpointCallbackConfig
+from .checkpoint import LastCheckpointCallback as LastCheckpointCallback
+from .checkpoint import LastCheckpointCallbackConfig as LastCheckpointCallbackConfig
+from .checkpoint import OnExceptionCheckpointCallback as OnExceptionCheckpointCallback
+from .checkpoint import (
+    OnExceptionCheckpointCallbackConfig as OnExceptionCheckpointCallbackConfig,
+)
+from .debug_flag import DebugFlagCallback as DebugFlagCallback
+from .debug_flag import DebugFlagCallbackConfig as DebugFlagCallbackConfig
+from .directory_setup import DirectorySetupCallback as DirectorySetupCallback
+from .directory_setup import (
+    DirectorySetupCallbackConfig as DirectorySetupCallbackConfig,
+)
+from .early_stopping import EarlyStoppingCallback as EarlyStoppingCallback
+from .early_stopping import EarlyStoppingCallbackConfig as EarlyStoppingCallbackConfig
+from .ema import EMACallback as EMACallback
+from .ema import EMACallbackConfig as EMACallbackConfig
+from .finite_checks import FiniteChecksCallback as FiniteChecksCallback
+from .finite_checks import FiniteChecksCallbackConfig as FiniteChecksCallbackConfig
+from .gradient_skipping import GradientSkippingCallback as GradientSkippingCallback
+from .gradient_skipping import (
+    GradientSkippingCallbackConfig as GradientSkippingCallbackConfig,
+)
+from .interval import EpochIntervalCallback as EpochIntervalCallback
+from .interval import IntervalCallback as IntervalCallback
+from .interval import StepIntervalCallback as StepIntervalCallback
+from .log_epoch import LogEpochCallback as LogEpochCallback
+from .log_epoch import LogEpochCallbackConfig as LogEpochCallbackConfig
+from .norm_logging import NormLoggingCallback as NormLoggingCallback
+from .norm_logging import NormLoggingCallbackConfig as NormLoggingCallbackConfig
+from .print_table import PrintTableMetricsCallback as PrintTableMetricsCallback
+from .print_table import (
+    PrintTableMetricsCallbackConfig as PrintTableMetricsCallbackConfig,
+)
+from .rlp_sanity_checks import RLPSanityChecksCallback as RLPSanityChecksCallback
+from .rlp_sanity_checks import (
+    RLPSanityChecksCallbackConfig as RLPSanityChecksCallbackConfig,
+)
+from .shared_parameters import SharedParametersCallback as SharedParametersCallback
+from .shared_parameters import (
+    SharedParametersCallbackConfig as SharedParametersCallbackConfig,
+)
+from .throughput_monitor import ThroughputMonitorConfig as ThroughputMonitorConfig
+from .timer import EpochTimerCallback as EpochTimerCallback
+from .timer import EpochTimerCallbackConfig as EpochTimerCallbackConfig
+from .wandb_upload_code import WandbUploadCodeCallback as WandbUploadCodeCallback
+from .wandb_upload_code import (
+    WandbUploadCodeCallbackConfig as WandbUploadCodeCallbackConfig,
+)
+from .wandb_watch import WandbWatchCallback as WandbWatchCallback
+from .wandb_watch import WandbWatchCallbackConfig as WandbWatchCallbackConfig
+
+CallbackConfig = Annotated[
+    DebugFlagCallbackConfig
+    | EarlyStoppingCallbackConfig
+    | ThroughputMonitorConfig
+    | EpochTimerCallbackConfig
+    | PrintTableMetricsCallbackConfig
+    | FiniteChecksCallbackConfig
+    | NormLoggingCallbackConfig
+    | GradientSkippingCallbackConfig
+    | LogEpochCallbackConfig
+    | EMACallbackConfig
+    | BestCheckpointCallbackConfig
+    | LastCheckpointCallbackConfig
+    | OnExceptionCheckpointCallbackConfig
+    | SharedParametersCallbackConfig
+    | RLPSanityChecksCallbackConfig
+    | WandbWatchCallbackConfig
+    | WandbUploadCodeCallbackConfig,
+    C.Field(discriminator="name"),
+]

{nshtrainer-0.41.1 → nshtrainer-0.43.0}/src/nshtrainer/callbacks/_throughput_monitor_callback.py

@@ -12,6 +12,8 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
+from __future__ import annotations
+
 import time
 from collections import deque
 from typing import (

{nshtrainer-0.41.1 → nshtrainer-0.43.0}/src/nshtrainer/callbacks/actsave.py

@@ -1,3 +1,5 @@
+from __future__ import annotations
+
 import contextlib
 from pathlib import Path
 from typing import Literal

{nshtrainer-0.41.1 → nshtrainer-0.43.0}/src/nshtrainer/callbacks/base.py

@@ -1,3 +1,5 @@
+from __future__ import annotations
+
 from abc import ABC, abstractmethod
 from collections import Counter
 from collections.abc import Iterable

{nshtrainer-0.41.1 → nshtrainer-0.43.0}/src/nshtrainer/callbacks/checkpoint/__init__.py

@@ -1,12 +1,16 @@
+from __future__ import annotations
+
 from .best_checkpoint import BestCheckpoint as BestCheckpoint
 from .best_checkpoint import (
     BestCheckpointCallbackConfig as BestCheckpointCallbackConfig,
 )
-from .last_checkpoint import LastCheckpoint as LastCheckpoint
+from .last_checkpoint import LastCheckpointCallback as LastCheckpointCallback
 from .last_checkpoint import (
     LastCheckpointCallbackConfig as LastCheckpointCallbackConfig,
 )
-from .on_exception_checkpoint import OnExceptionCheckpoint as OnExceptionCheckpoint
+from .on_exception_checkpoint import (
+    OnExceptionCheckpointCallback as OnExceptionCheckpointCallback,
+)
 from .on_exception_checkpoint import (
     OnExceptionCheckpointCallbackConfig as OnExceptionCheckpointCallbackConfig,
 )

{nshtrainer-0.41.1 → nshtrainer-0.43.0}/src/nshtrainer/callbacks/checkpoint/_base.py

@@ -1,3 +1,5 @@
+from __future__ import annotations
+
 import logging
 from abc import ABC, abstractmethod
 from pathlib import Path

{nshtrainer-0.41.1 → nshtrainer-0.43.0}/src/nshtrainer/callbacks/checkpoint/best_checkpoint.py

@@ -1,3 +1,5 @@
+from __future__ import annotations
+
 import logging
 from pathlib import Path
 from typing import Literal

{nshtrainer-0.41.1 → nshtrainer-0.43.0}/src/nshtrainer/callbacks/checkpoint/last_checkpoint.py

@@ -1,3 +1,5 @@
+from __future__ import annotations
+
 import logging
 from typing import Literal
 
@@ -17,11 +19,11 @@ class LastCheckpointCallbackConfig(BaseCheckpointCallbackConfig):
 
     @override
     def create_checkpoint(self, root_config, dirpath):
-        return LastCheckpoint(self, dirpath)
+        return LastCheckpointCallback(self, dirpath)
 
 
 @final
-class LastCheckpoint(CheckpointBase[LastCheckpointCallbackConfig]):
+class LastCheckpointCallback(CheckpointBase[LastCheckpointCallbackConfig]):
     @override
     def name(self):
         return "last"

{nshtrainer-0.41.1 → nshtrainer-0.43.0}/src/nshtrainer/callbacks/checkpoint/on_exception_checkpoint.py

@@ -1,3 +1,5 @@
+from __future__ import annotations
+
 import contextlib
 import datetime
 import logging
@@ -59,10 +61,12 @@ class OnExceptionCheckpointCallbackConfig(CallbackConfigBase):
 
         if not (filename := self.filename):
             filename = f"on_exception_{root_config.id}"
-        yield OnExceptionCheckpoint(self, dirpath=Path(dirpath), filename=filename)
+        yield OnExceptionCheckpointCallback(
+            self, dirpath=Path(dirpath), filename=filename
+        )
 
 
-class OnExceptionCheckpoint(_OnExceptionCheckpoint):
+class OnExceptionCheckpointCallback(_OnExceptionCheckpoint):
     @override
     def __init__(
         self,

{nshtrainer-0.41.1 → nshtrainer-0.43.0}/src/nshtrainer/callbacks/debug_flag.py

@@ -1,3 +1,5 @@
+from __future__ import annotations
+
 import logging
 from typing import TYPE_CHECKING, Literal, cast
 

{nshtrainer-0.41.1 → nshtrainer-0.43.0}/src/nshtrainer/callbacks/directory_setup.py

@@ -1,3 +1,5 @@
+from __future__ import annotations
+
 import logging
 import os
 from pathlib import Path
@@ -41,7 +43,7 @@ def _create_symlink_to_nshrunner(base_dir: Path):
     symlink_path.symlink_to(session_dir)
 
 
-class DirectorySetupConfig(CallbackConfigBase):
+class DirectorySetupCallbackConfig(CallbackConfigBase):
     name: Literal["directory_setup"] = "directory_setup"
 
     enabled: bool = True
@@ -62,7 +64,7 @@ class DirectorySetupConfig(CallbackConfigBase):
 
 class DirectorySetupCallback(Callback):
     @override
-    def __init__(self, config: DirectorySetupConfig):
+    def __init__(self, config: DirectorySetupCallbackConfig):
         super().__init__()
 
         self.config = config

{nshtrainer-0.41.1 → nshtrainer-0.43.0}/src/nshtrainer/callbacks/early_stopping.py

@@ -1,3 +1,5 @@
+from __future__ import annotations
+
 import logging
 import math
 from typing import Literal
@@ -14,7 +16,7 @@ from .base import CallbackConfigBase
 log = logging.getLogger(__name__)
 
 
-class EarlyStoppingConfig(CallbackConfigBase):
+class EarlyStoppingCallbackConfig(CallbackConfigBase):
     name: Literal["early_stopping"] = "early_stopping"
 
     metric: MetricConfig | None = None
@@ -54,11 +56,11 @@ class EarlyStoppingConfig(CallbackConfigBase):
                 "Either `metric` or `root_config.primary_metric` must be set to use EarlyStopping."
             )
 
-        yield EarlyStopping(self, metric)
+        yield EarlyStoppingCallback(self, metric)
 
 
-class EarlyStopping(_EarlyStopping):
-    def __init__(self, config: EarlyStoppingConfig, metric: MetricConfig):
+class EarlyStoppingCallback(_EarlyStopping):
+    def __init__(self, config: EarlyStoppingCallbackConfig, metric: MetricConfig):
         self.config = config
         self.metric = metric
         del config, metric

{nshtrainer-0.41.1 → nshtrainer-0.43.0}/src/nshtrainer/callbacks/ema.py

@@ -1,3 +1,5 @@
+from __future__ import annotations
+
 import contextlib
 import copy
 import threading
@@ -13,7 +15,7 @@ from typing_extensions import override
 from .base import CallbackConfigBase
 
 
-class EMA(Callback):
+class EMACallback(Callback):
     """
     Implements Exponential Moving Averaging (EMA).
 
@@ -358,7 +360,7 @@ class EMAOptimizer(torch.optim.Optimizer):
         self.rebuild_ema_params = True
 
 
-class EMAConfig(CallbackConfigBase):
+class EMACallbackConfig(CallbackConfigBase):
     name: Literal["ema"] = "ema"
 
     decay: float
@@ -375,7 +377,7 @@ class EMAConfig(CallbackConfigBase):
 
     @override
     def create_callbacks(self, root_config):
-        yield EMA(
+        yield EMACallback(
             decay=self.decay,
             validate_original_weights=self.validate_original_weights,
             every_n_steps=self.every_n_steps,

{nshtrainer-0.41.1 → nshtrainer-0.43.0}/src/nshtrainer/callbacks/finite_checks.py

@@ -1,3 +1,5 @@
+from __future__ import annotations
+
 import logging
 from typing import Literal
 
@@ -58,7 +60,7 @@ class FiniteChecksCallback(Callback):
         )
 
 
-class FiniteChecksConfig(CallbackConfigBase):
+class FiniteChecksCallbackConfig(CallbackConfigBase):
     name: Literal["finite_checks"] = "finite_checks"
 
     nonfinite_grads: bool = True

{nshtrainer-0.41.1 → nshtrainer-0.43.0}/src/nshtrainer/callbacks/gradient_skipping.py

@@ -1,3 +1,5 @@
+from __future__ import annotations
+
 import logging
 from typing import Any, Literal, Protocol, runtime_checkable
 
@@ -18,8 +20,8 @@ class HasGradSkippedSteps(Protocol):
     grad_skipped_steps: Any
 
 
-class GradientSkipping(Callback):
-    def __init__(self, config: "GradientSkippingConfig"):
+class GradientSkippingCallback(Callback):
+    def __init__(self, config: "GradientSkippingCallbackConfig"):
         super().__init__()
         self.config = config
 
@@ -73,7 +75,7 @@ class GradientSkipping(Callback):
         )
 
 
-class GradientSkippingConfig(CallbackConfigBase):
+class GradientSkippingCallbackConfig(CallbackConfigBase):
     name: Literal["gradient_skipping"] = "gradient_skipping"
 
     threshold: float
@@ -94,4 +96,4 @@ class GradientSkippingConfig(CallbackConfigBase):
 
     @override
     def create_callbacks(self, root_config):
-        yield GradientSkipping(self)
+        yield GradientSkippingCallback(self)

{nshtrainer-0.41.1 → nshtrainer-0.43.0}/src/nshtrainer/callbacks/interval.py

@@ -1,3 +1,5 @@
+from __future__ import annotations
+
 from collections.abc import Callable
 from typing import Literal
 

{nshtrainer-0.41.1 → nshtrainer-0.43.0}/src/nshtrainer/callbacks/log_epoch.py

@@ -1,14 +1,26 @@
+from __future__ import annotations
+
 import logging
 import math
-from typing import Any
+from typing import Any, Literal
 
 from lightning.pytorch import LightningModule, Trainer
 from lightning.pytorch.callbacks import Callback
 from typing_extensions import override
 
+from .base import CallbackConfigBase
+
 log = logging.getLogger(__name__)
 
 
+class LogEpochCallbackConfig(CallbackConfigBase):
+    name: Literal["log_epoch"] = "log_epoch"
+
+    @override
+    def create_callbacks(self, root_config):
+        yield LogEpochCallback()
+
+
 class LogEpochCallback(Callback):
     def __init__(self, metric_name: str = "computed_epoch"):
         super().__init__()

{nshtrainer-0.41.1 → nshtrainer-0.43.0}/src/nshtrainer/callbacks/norm_logging.py

@@ -1,3 +1,5 @@
+from __future__ import annotations
+
 import logging
 from typing import Literal, cast
 
@@ -96,7 +98,7 @@ def compute_norm(
 
 
 class NormLoggingCallback(Callback):
-    def __init__(self, config: "NormLoggingConfig"):
+    def __init__(self, config: "NormLoggingCallbackConfig"):
         super().__init__()
 
         self.config = config
@@ -155,7 +157,7 @@ class NormLoggingCallback(Callback):
                 )
 
 
-class NormLoggingConfig(CallbackConfigBase):
+class NormLoggingCallbackConfig(CallbackConfigBase):
     name: Literal["norm_logging"] = "norm_logging"
 
     log_grad_norm: bool | str | float = False

{nshtrainer-0.41.1 → nshtrainer-0.43.0}/src/nshtrainer/callbacks/print_table.py

@@ -1,3 +1,5 @@
+from __future__ import annotations
+
 import copy
 import fnmatch
 import importlib.util
@@ -74,7 +76,7 @@ class PrintTableMetricsCallback(Callback):
         return table
 
 
-class PrintTableMetricsConfig(CallbackConfigBase):
+class PrintTableMetricsCallbackConfig(CallbackConfigBase):
     """Configuration class for PrintTableMetricsCallback."""
 
     name: Literal["print_table_metrics"] = "print_table_metrics"

{nshtrainer-0.41.1 → nshtrainer-0.43.0}/src/nshtrainer/callbacks/rlp_sanity_checks.py

@@ -1,3 +1,5 @@
+from __future__ import annotations
+
 import logging
 from collections.abc import Mapping
 from typing import Literal, cast
@@ -16,7 +18,7 @@ from .base import CallbackConfigBase
 log = logging.getLogger(__name__)
 
 
-class RLPSanityChecksConfig(CallbackConfigBase):
+class RLPSanityChecksCallbackConfig(CallbackConfigBase):
     """
     If enabled, will do some sanity checks if the `ReduceLROnPlateau` scheduler is used:
         - If the ``interval`` is step, it makes sure that validation is called every ``frequency`` steps.
@@ -43,7 +45,7 @@ class RLPSanityChecksConfig(CallbackConfigBase):
 
 class RLPSanityChecksCallback(Callback):
     @override
-    def __init__(self, config: RLPSanityChecksConfig):
+    def __init__(self, config: RLPSanityChecksCallbackConfig):
         super().__init__()
 
         self.config = config

{nshtrainer-0.41.1 → nshtrainer-0.43.0}/src/nshtrainer/callbacks/shared_parameters.py

@@ -1,3 +1,5 @@
+from __future__ import annotations
+
 import logging
 from collections.abc import Iterable
 from typing import Literal, Protocol, TypeAlias, runtime_checkable
@@ -17,7 +19,7 @@ def _parameters_to_names(parameters: Iterable[nn.Parameter], model: nn.Module):
     return [mapping[id(p)] for p in parameters]
 
 
-class SharedParametersConfig(CallbackConfigBase):
+class SharedParametersCallbackConfig(CallbackConfigBase):
     """A callback that allows scaling the gradients of shared parameters that
     are registered in the ``self.shared_parameters`` list of the root module.
 
@@ -43,7 +45,7 @@ class ModuleWithSharedParameters(Protocol):
 
 class SharedParametersCallback(Callback):
     @override
-    def __init__(self, config: SharedParametersConfig):
+    def __init__(self, config: SharedParametersCallbackConfig):
         super().__init__()
 
         self.config = config

{nshtrainer-0.41.1 → nshtrainer-0.43.0}/src/nshtrainer/callbacks/throughput_monitor.py

@@ -1,3 +1,5 @@
+from __future__ import annotations
+
 import logging
 from typing import Any, Literal, Protocol, TypedDict, cast, runtime_checkable
 

{nshtrainer-0.41.1 → nshtrainer-0.43.0}/src/nshtrainer/callbacks/timer.py

@@ -1,3 +1,5 @@
+from __future__ import annotations
+
 import logging
 import time
 from typing import Any, Literal
@@ -12,7 +14,7 @@ from .base import CallbackConfigBase
 log = logging.getLogger(__name__)
 
 
-class EpochTimer(Callback):
+class EpochTimerCallback(Callback):
     def __init__(self):
         super().__init__()
 
@@ -149,9 +151,9 @@ class EpochTimer(Callback):
         self._total_batches = state_dict["total_batches"]
 
 
-class EpochTimerConfig(CallbackConfigBase):
+class EpochTimerCallbackConfig(CallbackConfigBase):
     name: Literal["epoch_timer"] = "epoch_timer"
 
     @override
     def create_callbacks(self, root_config):
-        yield EpochTimer()
+        yield EpochTimerCallback()

{nshtrainer-0.41.1 → nshtrainer-0.43.0}/src/nshtrainer/callbacks/wandb_upload_code.py

@@ -1,3 +1,5 @@
+from __future__ import annotations
+
 import logging
 import os
 from pathlib import Path
@@ -14,7 +16,7 @@ from .base import CallbackConfigBase
 log = logging.getLogger(__name__)
 
 
-class WandbUploadCodeConfig(CallbackConfigBase):
+class WandbUploadCodeCallbackConfig(CallbackConfigBase):
     name: Literal["wandb_upload_code"] = "wandb_upload_code"
 
     enabled: bool = True
@@ -32,7 +34,7 @@ class WandbUploadCodeConfig(CallbackConfigBase):
 
 
 class WandbUploadCodeCallback(Callback):
-    def __init__(self, config: WandbUploadCodeConfig):
+    def __init__(self, config: WandbUploadCodeCallbackConfig):
         super().__init__()
 
         self.config = config

{nshtrainer-0.41.1 → nshtrainer-0.43.0}/src/nshtrainer/callbacks/wandb_watch.py

@@ -1,3 +1,5 @@
+from __future__ import annotations
+
 import logging
 from typing import Literal, Protocol, cast, runtime_checkable
 
@@ -12,7 +14,7 @@ from .base import CallbackConfigBase
 log = logging.getLogger(__name__)
 
 
-class WandbWatchConfig(CallbackConfigBase):
+class WandbWatchCallbackConfig(CallbackConfigBase):
     name: Literal["wandb_watch"] = "wandb_watch"
 
     enabled: bool = True
@@ -41,7 +43,7 @@ class _HasWandbLogModuleProtocol(Protocol):
 
 
 class WandbWatchCallback(Callback):
-    def __init__(self, config: WandbWatchConfig):
+    def __init__(self, config: WandbWatchCallbackConfig):
         super().__init__()
 
         self.config = config