adding argument to predict()

Customized_IDSF/load_model.py

@@ -120,7 +120,7 @@ def convert_input_file_to_tensor_dataset(
 
     return dataset
 
-def predict(text, pred_config, args, tokenizer, pad_token_label_id):
+def predict(text, pred_config, args, tokenizer, pad_token_label_id, device):
 
     lines = text
     dataset = convert_input_file_to_tensor_dataset(lines, pred_config, args, tokenizer, pad_token_label_id)