diff --git a/declearn/training/_manager.py b/declearn/training/_manager.py index a3d158f8dcf3d0ccc4dae7ca8888cd6a1075111f..bfe21be65c3864a393cd519445dad8f5cdebe1c3 100644 --- a/declearn/training/_manager.py +++ b/declearn/training/_manager.py @@ -23,8 +23,8 @@ from typing import Any, Dict, List, Optional, Tuple, Union import numpy as np import tqdm +from declearn import messaging from declearn.aggregator import Aggregator -from declearn.communication import messaging from declearn.dataset import Dataset from declearn.metrics import ( MeanMetric, diff --git a/declearn/training/dp/_manager.py b/declearn/training/dp/_manager.py index 26aa2245a8e7f71b221a5da64ce00bec7ecd113a..291bea84caeee4eae16f2596f6f51b72b1da4ede 100644 --- a/declearn/training/dp/_manager.py +++ b/declearn/training/dp/_manager.py @@ -23,8 +23,8 @@ from typing import List, Optional, Tuple, Union from opacus.accountants import IAccountant, create_accountant # type: ignore from opacus.accountants.utils import get_noise_multiplier # type: ignore +from declearn import messaging from declearn.aggregator import Aggregator -from declearn.communication import messaging from declearn.dataset import Dataset from declearn.metrics import MetricInputType, MetricSet from declearn.model.api import Model diff --git a/test/training/test_train_manager.py b/test/training/test_train_manager.py index 271c2a1f52d1d8a88596f66315f2101c2f1403b2..f59ae5b154ed7bcc2457f6b5ce58aed7c7d876be 100644 --- a/test/training/test_train_manager.py +++ b/test/training/test_train_manager.py @@ -22,8 +22,8 @@ from typing import Any, Iterator, Optional import numpy +from declearn import messaging from declearn.aggregator import Aggregator -from declearn.communication import messaging from declearn.dataset import Dataset from declearn.metrics import Metric, MetricSet from declearn.model.api import Model, Vector diff --git a/test/training/test_train_manager_dp.py b/test/training/test_train_manager_dp.py index 01c62d7ff9c09f32d190055b6f147c94886d17d1..83edd26761b3738937afbdee326789d5b6db7847 100644 --- a/test/training/test_train_manager_dp.py +++ b/test/training/test_train_manager_dp.py @@ -28,7 +28,7 @@ try: except ModuleNotFoundError: pytest.skip("Opacus is unavailable", allow_module_level=True) -from declearn.communication import messaging +from declearn import messaging from declearn.dataset import DataSpecs from declearn.optimizer.modules import GaussianNoiseModule from declearn.training.dp import DPTrainingManager