aboutsummaryrefslogtreecommitdiff
path: root/src/mlia/nn/select.py
diff options
context:
space:
mode:
Diffstat (limited to 'src/mlia/nn/select.py')
-rw-r--r--src/mlia/nn/select.py4
1 files changed, 2 insertions, 2 deletions
diff --git a/src/mlia/nn/select.py b/src/mlia/nn/select.py
index 983426b..6947206 100644
--- a/src/mlia/nn/select.py
+++ b/src/mlia/nn/select.py
@@ -16,7 +16,7 @@ from mlia.core.errors import ConfigurationError
from mlia.nn.common import Optimizer
from mlia.nn.common import OptimizerConfiguration
from mlia.nn.rewrite.core.rewrite import RewriteConfiguration
-from mlia.nn.rewrite.core.rewrite import Rewriter
+from mlia.nn.rewrite.core.rewrite import RewritingOptimizer
from mlia.nn.rewrite.core.rewrite import TrainingParameters
from mlia.nn.tensorflow.config import KerasModel
from mlia.nn.tensorflow.config import TFLiteModel
@@ -132,7 +132,7 @@ def get_optimizer(
return Clusterer(model, config)
if isinstance(config, RewriteConfiguration):
- return Rewriter(model, config)
+ return RewritingOptimizer(model, config)
if isinstance(config, OptimizationSettings):
return _get_optimizer(model, cast(OptimizationSettings, config))