fix re inference error (PaddlePaddle#8473)

truongnn1106 · Dec 1, 2022 · 29a33c5 · 29a33c5
1 parent 7214949
commit 29a33c5
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/tools/infer_kie_token_ser_re.py b/tools/infer_kie_token_ser_re.py
@@ -81,7 +81,7 @@ def make_input(ser_inputs, ser_results):
  end.append(entity['end'])
  label.append(entities_labels[res['pred']])
 
- entities = np.full([max_seq_len + 1, 3], fill_value=-1)
+ entities = np.full([max_seq_len + 1, 3], fill_value=-1, dtype=np.int64)
  entities[0, 0] = len(start)
  entities[1:len(start) + 1, 0] = start
  entities[0, 1] = len(end)
@@ -98,7 +98,7 @@ def make_input(ser_inputs, ser_results):
  head.append(i)
  tail.append(j)
 
- relations = np.full([len(head) + 1, 2], fill_value=-1)
+ relations = np.full([len(head) + 1, 2], fill_value=-1, dtype=np.int64)
  relations[0, 0] = len(head)
  relations[1:len(head) + 1, 0] = head
  relations[0, 1] = len(tail)