devve1 commited on
Commit
2e66948
1 Parent(s): a9ba889

Update optimum_encoder.py

Browse files
Files changed (1) hide show
  1. 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
- providers=['TensorrtExecutionProvider'],
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