Update pairwise_model.py

pairwise_model.py

@@ -8,7 +8,7 @@ from optimum.intel import OVModelForQuestionAnswering
 import openvino.inference_engine as ie
 import os
 import gradio as gr
-
+from multiprocessing import cpu_count
 AUTH_TOKEN = "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"
 
 tokenizer = AutoTokenizer.from_pretrained('nguyenvulebinh/vi-mrc-base',
@@ -35,7 +35,12 @@ class PairwiseModel_modify(nn.Module):
 
     def forward(self, ids, masks):
         # Export the model to ONNX format
-        input_feed = {"input_ids": ids.cpu().numpy().astype(np.int64), "attention_mask": masks.cpu().numpy().astype(np.int64)}
+        ids_np = ids.cpu().numpy().astype(np.int64)
+        masks_np = masks.cpu().numpy().astype(np.int64)
+        ids_device = torch.from_numpy(ids_np).to(self.device)
+        masks_device = torch.from_numpy(masks_np).to(self.device)
+
+        input_feed = {"input_ids": ids_device, "attention_mask": masks_device}
         # Specify the input shapes (batch_size, max_sequence_length)
         input_shapes = {"input_ids": ids.shape, "attention_mask": masks.shape}
 
@@ -57,7 +62,7 @@ class PairwiseModel_modify(nn.Module):
         tmp["question"] = question
         valid_dataset = SiameseDatasetStage1(tmp, tokenizer, self.max_length, is_test=True)
         valid_loader = DataLoader(valid_dataset, batch_size=self.batch_size, collate_fn=collate_fn,
-                                  num_workers=0, shuffle=False, pin_memory=True)
+                                  num_workers=cpu_count(), shuffle=False, pin_memory=True)
         preds = []
         with torch.no_grad():
             bar = enumerate(valid_loader)