update trainer

weakly_supervised_parser/model/trainer.py

@@ -10,7 +10,7 @@ from pytorch_lightning import Trainer, seed_everything
 from pytorch_lightning.callbacks import EarlyStopping, ModelCheckpoint
 from transformers import AutoTokenizer, logging
 
-from onnxruntime import InferenceSession
+from onnxruntime import InferenceSession, SessionOptions
 from scipy.special import softmax
 
 from weakly_supervised_parser.model.data_module_loader import DataModule
@@ -98,7 +98,10 @@ class InsideOutsideStringClassifier:
         )
 
     def load_model(self, pre_trained_model_path):
-        self.model = InferenceSession(pre_trained_model_path, providers=["CUDAExecutionProvider", "CPUExecutionProvider"])
+        options = SessionOptions()
+        options.intra_op_num_threads = 1
+        options.inter_op_num_threads = 1
+        self.model = InferenceSession(pre_trained_model_path, options, providers=["CUDAExecutionProvider", "CPUExecutionProvider"])
         self.tokenizer = AutoTokenizer.from_pretrained(self.model_name_or_path, use_fast=True)
 
     def preprocess_function(self, data):