Update optimum_encoder.py
Browse files- optimum_encoder.py +5 -1
optimum_encoder.py
CHANGED
@@ -64,11 +64,15 @@ class OptimumEncoder(BaseEncoder):
|
|
64 |
"trt_fp16_enable": True
|
65 |
}
|
66 |
|
|
|
|
|
|
|
67 |
ort_model = ORTModelForFeatureExtraction.from_pretrained(
|
68 |
self.name,
|
69 |
model_save_dir=os.getenv('HF_HOME'),
|
70 |
-
|
71 |
provider_options=provider_options,
|
|
|
72 |
**self.model_kwargs
|
73 |
)
|
74 |
|
|
|
64 |
"trt_fp16_enable": True
|
65 |
}
|
66 |
|
67 |
+
session_options = onnxruntime.SessionOptions()
|
68 |
+
session_options.log_severity_level = 0
|
69 |
+
|
70 |
ort_model = ORTModelForFeatureExtraction.from_pretrained(
|
71 |
self.name,
|
72 |
model_save_dir=os.getenv('HF_HOME'),
|
73 |
+
provider=['TensorrtExecutionProvider'],
|
74 |
provider_options=provider_options,
|
75 |
+
session_options=session_options,
|
76 |
**self.model_kwargs
|
77 |
)
|
78 |
|