From d1236ddde64576024996abfbe07073519d2d99ef Mon Sep 17 00:00:00 2001
From: IlyasMoutawwakil <moutawwakil.ilyas.tsi@gmail.com>
Date: Thu, 24 Oct 2024 10:58:12 +0200
Subject: [PATCH] style

---
 optimum/exporters/onnx/model_patcher.py | 5 +----
 1 file changed, 1 insertion(+), 4 deletions(-)

diff --git a/optimum/exporters/onnx/model_patcher.py b/optimum/exporters/onnx/model_patcher.py
index cad4a41880..9b8119dc2f 100644
--- a/optimum/exporters/onnx/model_patcher.py
+++ b/optimum/exporters/onnx/model_patcher.py
@@ -1223,14 +1223,11 @@ def _prepare_4d_causal_attention_mask_with_cache_position_patched(
     return causal_mask
 
 
-from transformers import Phi3ForCausalLM
-
-
 class Phi3ModelPatcher(ModelPatcher):
     def __init__(
         self,
         config: "OnnxConfig",
-        model: Phi3ForCausalLM,
+        model: Union["PreTrainedModel", "TFPreTrainedModel"],
         model_kwargs: Optional[Dict[str, Any]] = None,
     ):
         super().__init__(config, model, model_kwargs)