training_utils: export neptune config and mixin

Like W&B. With this, #371 would have broken `test_import_training_utils`.

Follow up of #371 and #372.
This commit is contained in:
Cédric Deltheil 2024-05-29 07:40:18 +00:00 committed by Cédric Deltheil
parent 7dba1c8034
commit ddcd84c740

View file

@ -41,6 +41,7 @@ from refiners.training_utils.config import (
Optimizers, Optimizers,
TrainingConfig, TrainingConfig,
) )
from refiners.training_utils.neptune import NeptuneConfig, NeptuneMixin
from refiners.training_utils.trainer import Trainer, register_callback, register_model from refiners.training_utils.trainer import Trainer, register_callback, register_model
from refiners.training_utils.wandb import WandbConfig, WandbMixin from refiners.training_utils.wandb import WandbConfig, WandbMixin
@ -54,6 +55,8 @@ __all__ = [
"CallbackConfig", "CallbackConfig",
"WandbMixin", "WandbMixin",
"WandbConfig", "WandbConfig",
"NeptuneMixin",
"NeptuneConfig",
"LRSchedulerConfig", "LRSchedulerConfig",
"OptimizerConfig", "OptimizerConfig",
"TrainingConfig", "TrainingConfig",