From 734c6054d3a045c33881994837d7fdacc4d8d59a Mon Sep 17 00:00:00 2001 From: Jlonge4 Date: Sat, 1 Feb 2025 10:34:15 -0500 Subject: [PATCH] lint --- optimum/exporters/neuron/model_configs/decoder_configs.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/optimum/exporters/neuron/model_configs/decoder_configs.py b/optimum/exporters/neuron/model_configs/decoder_configs.py index d09abbf0a..10708024d 100644 --- a/optimum/exporters/neuron/model_configs/decoder_configs.py +++ b/optimum/exporters/neuron/model_configs/decoder_configs.py @@ -14,7 +14,6 @@ # limitations under the License. """Neuron export configurations for models using transformers_neuronx.""" - from optimum.exporters.tasks import TasksManager from ....neuron.models.granite.model import GraniteForSampling @@ -72,6 +71,7 @@ class GraniteNeuronConfig(TextNeuronDecoderConfig): NEURONX_CLASS = GraniteForSampling CONTINUOUS_BATCHING = True + @register_in_tasks_manager("phi4", "text-generation") class Phi4NeuronConfig(TextNeuronDecoderConfig): NEURONX_CLASS = Phi4ForSampling