style

huggingface · Oct 24, 2024 · d1236dd · d1236dd
1 parent 72ca53c
commit d1236dd
Showing 1 changed file with 1 addition and 4 deletions.
diff --git 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)