nshtrainer 0.42.0__py3-none-any.whl → 0.44.0__py3-none-any.whl

nshtrainer/config/callbacks/actsave/__init__.py

@@ -1,3 +1,5 @@
+from __future__ import annotations
+
 __codegen__ = True
 
 from typing import TYPE_CHECKING
@@ -14,12 +16,12 @@ else:
 
         if name in globals():
             return globals()[name]
+        if name == "ActSaveConfig":
+            return importlib.import_module("nshtrainer.callbacks.actsave").ActSaveConfig
         if name == "CallbackConfigBase":
             return importlib.import_module(
                 "nshtrainer.callbacks.actsave"
             ).CallbackConfigBase
-        if name == "ActSaveConfig":
-            return importlib.import_module("nshtrainer.callbacks.actsave").ActSaveConfig
         raise AttributeError(f"module '{__name__}' has no attribute '{name}'")
 
 # Submodule exports

nshtrainer/config/callbacks/base/__init__.py

@@ -1,3 +1,5 @@
+from __future__ import annotations
+
 __codegen__ = True
 
 from typing import TYPE_CHECKING

nshtrainer/config/callbacks/checkpoint/__init__.py

@@ -1,3 +1,5 @@
+from __future__ import annotations
+
 __codegen__ = True
 
 from typing import TYPE_CHECKING
@@ -33,10 +35,6 @@ else:
 
         if name in globals():
             return globals()[name]
-        if name == "LastCheckpointCallbackConfig":
-            return importlib.import_module(
-                "nshtrainer.callbacks.checkpoint"
-            ).LastCheckpointCallbackConfig
         if name == "CheckpointMetadata":
             return importlib.import_module(
                 "nshtrainer.callbacks.checkpoint._base"
@@ -45,6 +43,10 @@ else:
             return importlib.import_module(
                 "nshtrainer.callbacks.checkpoint._base"
             ).BaseCheckpointCallbackConfig
+        if name == "LastCheckpointCallbackConfig":
+            return importlib.import_module(
+                "nshtrainer.callbacks.checkpoint"
+            ).LastCheckpointCallbackConfig
         if name == "CallbackConfigBase":
             return importlib.import_module(
                 "nshtrainer.callbacks.checkpoint._base"

nshtrainer/config/callbacks/checkpoint/_base/__init__.py

@@ -1,3 +1,5 @@
+from __future__ import annotations
+
 __codegen__ = True
 
 from typing import TYPE_CHECKING
@@ -21,10 +23,6 @@ else:
 
         if name in globals():
             return globals()[name]
-        if name == "CallbackConfigBase":
-            return importlib.import_module(
-                "nshtrainer.callbacks.checkpoint._base"
-            ).CallbackConfigBase
         if name == "CheckpointMetadata":
             return importlib.import_module(
                 "nshtrainer.callbacks.checkpoint._base"
@@ -33,6 +31,10 @@ else:
             return importlib.import_module(
                 "nshtrainer.callbacks.checkpoint._base"
             ).BaseCheckpointCallbackConfig
+        if name == "CallbackConfigBase":
+            return importlib.import_module(
+                "nshtrainer.callbacks.checkpoint._base"
+            ).CallbackConfigBase
         raise AttributeError(f"module '{__name__}' has no attribute '{name}'")
 
 # Submodule exports

nshtrainer/config/callbacks/checkpoint/best_checkpoint/__init__.py

@@ -1,3 +1,5 @@
+from __future__ import annotations
+
 __codegen__ = True
 
 from typing import TYPE_CHECKING

nshtrainer/config/callbacks/checkpoint/last_checkpoint/__init__.py

@@ -1,3 +1,5 @@
+from __future__ import annotations
+
 __codegen__ = True
 
 from typing import TYPE_CHECKING
@@ -21,10 +23,6 @@ else:
 
         if name in globals():
             return globals()[name]
-        if name == "LastCheckpointCallbackConfig":
-            return importlib.import_module(
-                "nshtrainer.callbacks.checkpoint.last_checkpoint"
-            ).LastCheckpointCallbackConfig
         if name == "CheckpointMetadata":
             return importlib.import_module(
                 "nshtrainer.callbacks.checkpoint.last_checkpoint"
@@ -33,6 +31,10 @@ else:
             return importlib.import_module(
                 "nshtrainer.callbacks.checkpoint.last_checkpoint"
             ).BaseCheckpointCallbackConfig
+        if name == "LastCheckpointCallbackConfig":
+            return importlib.import_module(
+                "nshtrainer.callbacks.checkpoint.last_checkpoint"
+            ).LastCheckpointCallbackConfig
         raise AttributeError(f"module '{__name__}' has no attribute '{name}'")
 
 # Submodule exports

nshtrainer/config/callbacks/checkpoint/on_exception_checkpoint/__init__.py

@@ -1,3 +1,5 @@
+from __future__ import annotations
+
 __codegen__ = True
 
 from typing import TYPE_CHECKING
@@ -18,14 +20,14 @@ else:
 
         if name in globals():
             return globals()[name]
-        if name == "CallbackConfigBase":
-            return importlib.import_module(
-                "nshtrainer.callbacks.checkpoint.on_exception_checkpoint"
-            ).CallbackConfigBase
         if name == "OnExceptionCheckpointCallbackConfig":
             return importlib.import_module(
                 "nshtrainer.callbacks.checkpoint.on_exception_checkpoint"
             ).OnExceptionCheckpointCallbackConfig
+        if name == "CallbackConfigBase":
+            return importlib.import_module(
+                "nshtrainer.callbacks.checkpoint.on_exception_checkpoint"
+            ).CallbackConfigBase
         raise AttributeError(f"module '{__name__}' has no attribute '{name}'")
 
 # Submodule exports

nshtrainer/config/callbacks/debug_flag/__init__.py

@@ -1,3 +1,5 @@
+from __future__ import annotations
+
 __codegen__ = True
 
 from typing import TYPE_CHECKING
@@ -16,14 +18,14 @@ else:
 
         if name in globals():
             return globals()[name]
-        if name == "CallbackConfigBase":
-            return importlib.import_module(
-                "nshtrainer.callbacks.debug_flag"
-            ).CallbackConfigBase
         if name == "DebugFlagCallbackConfig":
             return importlib.import_module(
                 "nshtrainer.callbacks.debug_flag"
             ).DebugFlagCallbackConfig
+        if name == "CallbackConfigBase":
+            return importlib.import_module(
+                "nshtrainer.callbacks.debug_flag"
+            ).CallbackConfigBase
         raise AttributeError(f"module '{__name__}' has no attribute '{name}'")
 
 # Submodule exports

nshtrainer/config/callbacks/directory_setup/__init__.py

@@ -1,3 +1,5 @@
+from __future__ import annotations
+
 __codegen__ = True
 
 from typing import TYPE_CHECKING
@@ -9,7 +11,7 @@ if TYPE_CHECKING:
         CallbackConfigBase as CallbackConfigBase,
     )
     from nshtrainer.callbacks.directory_setup import (
-        DirectorySetupConfig as DirectorySetupConfig,
+        DirectorySetupCallbackConfig as DirectorySetupCallbackConfig,
     )
 else:
 
@@ -18,14 +20,14 @@ else:
 
         if name in globals():
             return globals()[name]
-        if name == "CallbackConfigBase":
+        if name == "DirectorySetupCallbackConfig":
             return importlib.import_module(
                 "nshtrainer.callbacks.directory_setup"
-            ).CallbackConfigBase
-        if name == "DirectorySetupConfig":
+            ).DirectorySetupCallbackConfig
+        if name == "CallbackConfigBase":
             return importlib.import_module(
                 "nshtrainer.callbacks.directory_setup"
-            ).DirectorySetupConfig
+            ).CallbackConfigBase
         raise AttributeError(f"module '{__name__}' has no attribute '{name}'")
 
 # Submodule exports

nshtrainer/config/callbacks/early_stopping/__init__.py

@@ -1,3 +1,5 @@
+from __future__ import annotations
+
 __codegen__ = True
 
 from typing import TYPE_CHECKING
@@ -9,7 +11,7 @@ if TYPE_CHECKING:
         CallbackConfigBase as CallbackConfigBase,
     )
     from nshtrainer.callbacks.early_stopping import (
-        EarlyStoppingConfig as EarlyStoppingConfig,
+        EarlyStoppingCallbackConfig as EarlyStoppingCallbackConfig,
     )
     from nshtrainer.callbacks.early_stopping import MetricConfig as MetricConfig
 else:
@@ -19,18 +21,18 @@ else:
 
         if name in globals():
             return globals()[name]
-        if name == "CallbackConfigBase":
-            return importlib.import_module(
-                "nshtrainer.callbacks.early_stopping"
-            ).CallbackConfigBase
         if name == "MetricConfig":
             return importlib.import_module(
                 "nshtrainer.callbacks.early_stopping"
             ).MetricConfig
-        if name == "EarlyStoppingConfig":
+        if name == "EarlyStoppingCallbackConfig":
             return importlib.import_module(
                 "nshtrainer.callbacks.early_stopping"
-            ).EarlyStoppingConfig
+            ).EarlyStoppingCallbackConfig
+        if name == "CallbackConfigBase":
+            return importlib.import_module(
+                "nshtrainer.callbacks.early_stopping"
+            ).CallbackConfigBase
         raise AttributeError(f"module '{__name__}' has no attribute '{name}'")
 
 # Submodule exports

nshtrainer/config/callbacks/ema/__init__.py

@@ -1,3 +1,5 @@
+from __future__ import annotations
+
 __codegen__ = True
 
 from typing import TYPE_CHECKING
@@ -6,7 +8,7 @@ from typing import TYPE_CHECKING
 
 if TYPE_CHECKING:
     from nshtrainer.callbacks.ema import CallbackConfigBase as CallbackConfigBase
-    from nshtrainer.callbacks.ema import EMAConfig as EMAConfig
+    from nshtrainer.callbacks.ema import EMACallbackConfig as EMACallbackConfig
 else:
 
     def __getattr__(name):
@@ -14,12 +16,12 @@ else:
 
         if name in globals():
             return globals()[name]
+        if name == "EMACallbackConfig":
+            return importlib.import_module("nshtrainer.callbacks.ema").EMACallbackConfig
         if name == "CallbackConfigBase":
             return importlib.import_module(
                 "nshtrainer.callbacks.ema"
             ).CallbackConfigBase
-        if name == "EMAConfig":
-            return importlib.import_module("nshtrainer.callbacks.ema").EMAConfig
         raise AttributeError(f"module '{__name__}' has no attribute '{name}'")
 
 # Submodule exports

nshtrainer/config/callbacks/finite_checks/__init__.py

@@ -1,3 +1,5 @@
+from __future__ import annotations
+
 __codegen__ = True
 
 from typing import TYPE_CHECKING
@@ -9,7 +11,7 @@ if TYPE_CHECKING:
         CallbackConfigBase as CallbackConfigBase,
     )
     from nshtrainer.callbacks.finite_checks import (
-        FiniteChecksConfig as FiniteChecksConfig,
+        FiniteChecksCallbackConfig as FiniteChecksCallbackConfig,
     )
 else:
 
@@ -18,14 +20,14 @@ else:
 
         if name in globals():
             return globals()[name]
-        if name == "CallbackConfigBase":
+        if name == "FiniteChecksCallbackConfig":
             return importlib.import_module(
                 "nshtrainer.callbacks.finite_checks"
-            ).CallbackConfigBase
-        if name == "FiniteChecksConfig":
+            ).FiniteChecksCallbackConfig
+        if name == "CallbackConfigBase":
             return importlib.import_module(
                 "nshtrainer.callbacks.finite_checks"
-            ).FiniteChecksConfig
+            ).CallbackConfigBase
         raise AttributeError(f"module '{__name__}' has no attribute '{name}'")
 
 # Submodule exports

nshtrainer/config/callbacks/gradient_skipping/__init__.py

@@ -1,3 +1,5 @@
+from __future__ import annotations
+
 __codegen__ = True
 
 from typing import TYPE_CHECKING
@@ -9,7 +11,7 @@ if TYPE_CHECKING:
         CallbackConfigBase as CallbackConfigBase,
     )
     from nshtrainer.callbacks.gradient_skipping import (
-        GradientSkippingConfig as GradientSkippingConfig,
+        GradientSkippingCallbackConfig as GradientSkippingCallbackConfig,
     )
 else:
 
@@ -18,14 +20,14 @@ else:
 
         if name in globals():
             return globals()[name]
-        if name == "CallbackConfigBase":
+        if name == "GradientSkippingCallbackConfig":
             return importlib.import_module(
                 "nshtrainer.callbacks.gradient_skipping"
-            ).CallbackConfigBase
-        if name == "GradientSkippingConfig":
+            ).GradientSkippingCallbackConfig
+        if name == "CallbackConfigBase":
             return importlib.import_module(
                 "nshtrainer.callbacks.gradient_skipping"
-            ).GradientSkippingConfig
+            ).CallbackConfigBase
         raise AttributeError(f"module '{__name__}' has no attribute '{name}'")
 
 # Submodule exports

nshtrainer/config/callbacks/norm_logging/__init__.py

@@ -1,3 +1,5 @@
+from __future__ import annotations
+
 __codegen__ = True
 
 from typing import TYPE_CHECKING
@@ -8,7 +10,9 @@ if TYPE_CHECKING:
     from nshtrainer.callbacks.norm_logging import (
         CallbackConfigBase as CallbackConfigBase,
     )
-    from nshtrainer.callbacks.norm_logging import NormLoggingConfig as NormLoggingConfig
+    from nshtrainer.callbacks.norm_logging import (
+        NormLoggingCallbackConfig as NormLoggingCallbackConfig,
+    )
 else:
 
     def __getattr__(name):
@@ -16,14 +20,14 @@ else:
 
         if name in globals():
             return globals()[name]
-        if name == "CallbackConfigBase":
+        if name == "NormLoggingCallbackConfig":
             return importlib.import_module(
                 "nshtrainer.callbacks.norm_logging"
-            ).CallbackConfigBase
-        if name == "NormLoggingConfig":
+            ).NormLoggingCallbackConfig
+        if name == "CallbackConfigBase":
             return importlib.import_module(
                 "nshtrainer.callbacks.norm_logging"
-            ).NormLoggingConfig
+            ).CallbackConfigBase
         raise AttributeError(f"module '{__name__}' has no attribute '{name}'")
 
 # Submodule exports

nshtrainer/config/callbacks/print_table/__init__.py

@@ -1,3 +1,5 @@
+from __future__ import annotations
+
 __codegen__ = True
 
 from typing import TYPE_CHECKING
@@ -9,7 +11,7 @@ if TYPE_CHECKING:
         CallbackConfigBase as CallbackConfigBase,
     )
     from nshtrainer.callbacks.print_table import (
-        PrintTableMetricsConfig as PrintTableMetricsConfig,
+        PrintTableMetricsCallbackConfig as PrintTableMetricsCallbackConfig,
     )
 else:
 
@@ -18,14 +20,14 @@ else:
 
         if name in globals():
             return globals()[name]
-        if name == "CallbackConfigBase":
+        if name == "PrintTableMetricsCallbackConfig":
             return importlib.import_module(
                 "nshtrainer.callbacks.print_table"
-            ).CallbackConfigBase
-        if name == "PrintTableMetricsConfig":
+            ).PrintTableMetricsCallbackConfig
+        if name == "CallbackConfigBase":
             return importlib.import_module(
                 "nshtrainer.callbacks.print_table"
-            ).PrintTableMetricsConfig
+            ).CallbackConfigBase
         raise AttributeError(f"module '{__name__}' has no attribute '{name}'")
 
 # Submodule exports

nshtrainer/config/callbacks/rlp_sanity_checks/__init__.py

@@ -1,3 +1,5 @@
+from __future__ import annotations
+
 __codegen__ = True
 
 from typing import TYPE_CHECKING
@@ -9,7 +11,7 @@ if TYPE_CHECKING:
         CallbackConfigBase as CallbackConfigBase,
     )
     from nshtrainer.callbacks.rlp_sanity_checks import (
-        RLPSanityChecksConfig as RLPSanityChecksConfig,
+        RLPSanityChecksCallbackConfig as RLPSanityChecksCallbackConfig,
     )
 else:
 
@@ -18,14 +20,14 @@ else:
 
         if name in globals():
             return globals()[name]
-        if name == "CallbackConfigBase":
+        if name == "RLPSanityChecksCallbackConfig":
             return importlib.import_module(
                 "nshtrainer.callbacks.rlp_sanity_checks"
-            ).CallbackConfigBase
-        if name == "RLPSanityChecksConfig":
+            ).RLPSanityChecksCallbackConfig
+        if name == "CallbackConfigBase":
             return importlib.import_module(
                 "nshtrainer.callbacks.rlp_sanity_checks"
-            ).RLPSanityChecksConfig
+            ).CallbackConfigBase
         raise AttributeError(f"module '{__name__}' has no attribute '{name}'")
 
 # Submodule exports

nshtrainer/config/callbacks/shared_parameters/__init__.py

@@ -1,3 +1,5 @@
+from __future__ import annotations
+
 __codegen__ = True
 
 from typing import TYPE_CHECKING
@@ -9,7 +11,7 @@ if TYPE_CHECKING:
         CallbackConfigBase as CallbackConfigBase,
     )
     from nshtrainer.callbacks.shared_parameters import (
-        SharedParametersConfig as SharedParametersConfig,
+        SharedParametersCallbackConfig as SharedParametersCallbackConfig,
     )
 else:
 
@@ -18,14 +20,14 @@ else:
 
         if name in globals():
             return globals()[name]
-        if name == "CallbackConfigBase":
+        if name == "SharedParametersCallbackConfig":
             return importlib.import_module(
                 "nshtrainer.callbacks.shared_parameters"
-            ).CallbackConfigBase
-        if name == "SharedParametersConfig":
+            ).SharedParametersCallbackConfig
+        if name == "CallbackConfigBase":
             return importlib.import_module(
                 "nshtrainer.callbacks.shared_parameters"
-            ).SharedParametersConfig
+            ).CallbackConfigBase
         raise AttributeError(f"module '{__name__}' has no attribute '{name}'")
 
 # Submodule exports

nshtrainer/config/callbacks/throughput_monitor/__init__.py

@@ -1,3 +1,5 @@
+from __future__ import annotations
+
 __codegen__ = True
 
 from typing import TYPE_CHECKING
@@ -18,14 +20,14 @@ else:
 
         if name in globals():
             return globals()[name]
-        if name == "CallbackConfigBase":
-            return importlib.import_module(
-                "nshtrainer.callbacks.throughput_monitor"
-            ).CallbackConfigBase
         if name == "ThroughputMonitorConfig":
             return importlib.import_module(
                 "nshtrainer.callbacks.throughput_monitor"
             ).ThroughputMonitorConfig
+        if name == "CallbackConfigBase":
+            return importlib.import_module(
+                "nshtrainer.callbacks.throughput_monitor"
+            ).CallbackConfigBase
         raise AttributeError(f"module '{__name__}' has no attribute '{name}'")
 
 # Submodule exports

nshtrainer/config/callbacks/timer/__init__.py

@@ -1,3 +1,5 @@
+from __future__ import annotations
+
 __codegen__ = True
 
 from typing import TYPE_CHECKING
@@ -6,7 +8,9 @@ from typing import TYPE_CHECKING
 
 if TYPE_CHECKING:
     from nshtrainer.callbacks.timer import CallbackConfigBase as CallbackConfigBase
-    from nshtrainer.callbacks.timer import EpochTimerConfig as EpochTimerConfig
+    from nshtrainer.callbacks.timer import (
+        EpochTimerCallbackConfig as EpochTimerCallbackConfig,
+    )
 else:
 
     def __getattr__(name):
@@ -14,14 +18,14 @@ else:
 
         if name in globals():
             return globals()[name]
-        if name == "CallbackConfigBase":
+        if name == "EpochTimerCallbackConfig":
             return importlib.import_module(
                 "nshtrainer.callbacks.timer"
-            ).CallbackConfigBase
-        if name == "EpochTimerConfig":
+            ).EpochTimerCallbackConfig
+        if name == "CallbackConfigBase":
             return importlib.import_module(
                 "nshtrainer.callbacks.timer"
-            ).EpochTimerConfig
+            ).CallbackConfigBase
         raise AttributeError(f"module '{__name__}' has no attribute '{name}'")
 
 # Submodule exports

nshtrainer/config/callbacks/wandb_upload_code/__init__.py

@@ -1,3 +1,5 @@
+from __future__ import annotations
+
 __codegen__ = True
 
 from typing import TYPE_CHECKING
@@ -9,7 +11,7 @@ if TYPE_CHECKING:
         CallbackConfigBase as CallbackConfigBase,
     )
     from nshtrainer.callbacks.wandb_upload_code import (
-        WandbUploadCodeConfig as WandbUploadCodeConfig,
+        WandbUploadCodeCallbackConfig as WandbUploadCodeCallbackConfig,
     )
 else:
 
@@ -18,14 +20,14 @@ else:
 
         if name in globals():
             return globals()[name]
-        if name == "CallbackConfigBase":
+        if name == "WandbUploadCodeCallbackConfig":
             return importlib.import_module(
                 "nshtrainer.callbacks.wandb_upload_code"
-            ).CallbackConfigBase
-        if name == "WandbUploadCodeConfig":
+            ).WandbUploadCodeCallbackConfig
+        if name == "CallbackConfigBase":
             return importlib.import_module(
                 "nshtrainer.callbacks.wandb_upload_code"
-            ).WandbUploadCodeConfig
+            ).CallbackConfigBase
         raise AttributeError(f"module '{__name__}' has no attribute '{name}'")
 
 # Submodule exports

nshtrainer/config/callbacks/wandb_watch/__init__.py

@@ -1,3 +1,5 @@
+from __future__ import annotations
+
 __codegen__ = True
 
 from typing import TYPE_CHECKING
@@ -8,7 +10,9 @@ if TYPE_CHECKING:
     from nshtrainer.callbacks.wandb_watch import (
         CallbackConfigBase as CallbackConfigBase,
     )
-    from nshtrainer.callbacks.wandb_watch import WandbWatchConfig as WandbWatchConfig
+    from nshtrainer.callbacks.wandb_watch import (
+        WandbWatchCallbackConfig as WandbWatchCallbackConfig,
+    )
 else:
 
     def __getattr__(name):
@@ -16,14 +20,14 @@ else:
 
         if name in globals():
             return globals()[name]
-        if name == "CallbackConfigBase":
+        if name == "WandbWatchCallbackConfig":
             return importlib.import_module(
                 "nshtrainer.callbacks.wandb_watch"
-            ).CallbackConfigBase
-        if name == "WandbWatchConfig":
+            ).WandbWatchCallbackConfig
+        if name == "CallbackConfigBase":
             return importlib.import_module(
                 "nshtrainer.callbacks.wandb_watch"
-            ).WandbWatchConfig
+            ).CallbackConfigBase
         raise AttributeError(f"module '{__name__}' has no attribute '{name}'")
 
 # Submodule exports

nshtrainer/config/loggers/__init__.py

@@ -1,3 +1,5 @@
+from __future__ import annotations
+
 __codegen__ = True
 
 from typing import TYPE_CHECKING
@@ -11,8 +13,12 @@ if TYPE_CHECKING:
     from nshtrainer.loggers import TensorboardLoggerConfig as TensorboardLoggerConfig
     from nshtrainer.loggers import WandbLoggerConfig as WandbLoggerConfig
     from nshtrainer.loggers.wandb import CallbackConfigBase as CallbackConfigBase
-    from nshtrainer.loggers.wandb import WandbUploadCodeConfig as WandbUploadCodeConfig
-    from nshtrainer.loggers.wandb import WandbWatchConfig as WandbWatchConfig
+    from nshtrainer.loggers.wandb import (
+        WandbUploadCodeCallbackConfig as WandbUploadCodeCallbackConfig,
+    )
+    from nshtrainer.loggers.wandb import (
+        WandbWatchCallbackConfig as WandbWatchCallbackConfig,
+    )
 else:
 
     def __getattr__(name):
@@ -24,18 +30,20 @@ else:
             return importlib.import_module("nshtrainer.loggers").BaseLoggerConfig
         if name == "TensorboardLoggerConfig":
             return importlib.import_module("nshtrainer.loggers").TensorboardLoggerConfig
-        if name == "CallbackConfigBase":
-            return importlib.import_module(
-                "nshtrainer.loggers.wandb"
-            ).CallbackConfigBase
         if name == "WandbLoggerConfig":
             return importlib.import_module("nshtrainer.loggers").WandbLoggerConfig
-        if name == "WandbUploadCodeConfig":
+        if name == "WandbUploadCodeCallbackConfig":
             return importlib.import_module(
                 "nshtrainer.loggers.wandb"
-            ).WandbUploadCodeConfig
-        if name == "WandbWatchConfig":
-            return importlib.import_module("nshtrainer.loggers.wandb").WandbWatchConfig
+            ).WandbUploadCodeCallbackConfig
+        if name == "WandbWatchCallbackConfig":
+            return importlib.import_module(
+                "nshtrainer.loggers.wandb"
+            ).WandbWatchCallbackConfig
+        if name == "CallbackConfigBase":
+            return importlib.import_module(
+                "nshtrainer.loggers.wandb"
+            ).CallbackConfigBase
         if name == "CSVLoggerConfig":
             return importlib.import_module("nshtrainer.loggers").CSVLoggerConfig
         if name == "LoggerConfig":

nshtrainer/config/loggers/_base/__init__.py

@@ -1,3 +1,5 @@
+from __future__ import annotations
+
 __codegen__ = True
 
 from typing import TYPE_CHECKING