diff --git a/examples/onnxruntime/training/image-classification/run_image_classification.py b/examples/onnxruntime/training/image-classification/run_image_classification.py
index bc5482e23e..c5d5aabe27 100644
--- a/examples/onnxruntime/training/image-classification/run_image_classification.py
+++ b/examples/onnxruntime/training/image-classification/run_image_classification.py
@@ -48,7 +48,7 @@
 from transformers.utils import check_min_version, send_example_telemetry
 from transformers.utils.versions import require_version
 
-from optimum import ORTTrainer, ORTTrainingArguments
+from optimum.onnxruntime import ORTTrainer, ORTTrainingArguments
 
 
 """ Fine-tuning a 🤗 Transformers model for image classification"""