diff --git a/paddle/operators/beam_search_op.cc b/paddle/operators/beam_search_op.cc index 2e0513b37a24b..ed2e7b738acd8 100644 --- a/paddle/operators/beam_search_op.cc +++ b/paddle/operators/beam_search_op.cc @@ -39,7 +39,7 @@ void BeamSearch::operator()(const framework::LoDTensor &pre_ids, std::map> hash; framework::LoD new_lod; - auto *ids_data = selected_ids->mutable_data(platform::CPUPlace()); + auto *ids_data = selected_ids->mutable_data(platform::CPUPlace()); auto *scores_data = selected_scores->mutable_data(platform::CPUPlace()); @@ -66,7 +66,7 @@ void BeamSearch::operator()(const framework::LoDTensor &pre_ids, void BeamSearch::PruneEndidCandidates(const framework::LoDTensor &pre_ids, std::vector> *items) { - auto *pre_ids_data = pre_ids.data(); + auto *pre_ids_data = pre_ids.data(); for (size_t offset = 0; offset < items->size(); offset++) { auto prefix_id = pre_ids_data[offset]; @@ -127,7 +127,7 @@ bool BeamSearch::NextItemSet(std::vector *items) { auto abs_lod = framework::ToAbsOffset(ids.lod()); PADDLE_ENFORCE_GE(source_abs_two_level_lod.size(), 2UL); - auto *ids_data = ids.data(); + auto *ids_data = ids.data(); auto *scores_data = scores.data(); size_t instance_dim = 1; diff --git a/python/paddle/v2/fluid/tests/test_beam_search_op.py b/python/paddle/v2/fluid/tests/test_beam_search_op.py index 595f132fa85f0..319a7e49e35b0 100644 --- a/python/paddle/v2/fluid/tests/test_beam_search_op.py +++ b/python/paddle/v2/fluid/tests/test_beam_search_op.py @@ -37,13 +37,13 @@ def test_run(self): print 'lod', selected_ids.lod() def _create_pre_ids(self): - np_data = np.array([[1, 2, 3, 4]], dtype='int32') + np_data = np.array([[1, 2, 3, 4]], dtype='int64') tensor = create_tensor(self.scope, "pre_ids", np_data) def _create_ids(self): self.lod = [[0, 1, 4], [0, 1, 2, 3, 4]] np_data = np.array( - [[4, 2, 5], [2, 1, 3], [3, 5, 2], [8, 2, 1]], dtype='int32') + [[4, 2, 5], [2, 1, 3], [3, 5, 2], [8, 2, 1]], dtype='int64') tensor = create_tensor(self.scope, "ids", np_data) tensor.set_lod(self.lod)