fix batch_size type

supercoderhawk · supercoderhawk · commit 253c41b230ee · 2020-07-15T15:48:48.000+08:00
diff --git a/deep_keyphrase/copy_rnn/predict_tf.py b/deep_keyphrase/copy_rnn/predict_tf.py
@@ -29,7 +29,8 @@ def eval_predict(self, batch, model=None, delimiter=None):
     def predict(self, text_list):
         x_batch, x_oov_batch, sent_len_batch, oov_list_batch = self.generate_input_batch(text_list)
         batch_size = len(x_batch)
-        result_tensor = self.model.beam_search(x_batch, x_oov_batch, sent_len_batch, [batch_size])
+        batch_size_np = np.array([batch_size], dtype=np.long)
+        result_tensor = self.model.beam_search(x_batch, x_oov_batch, sent_len_batch, batch_size_np)
         result_np = result_tensor.numpy()
         return self.__idx2result_beam('', oov_list_batch, result_np)
 
