@@ -65,19 +65,26 @@ class FilterByInstagKernel : public framework::OpKernel<T> {
6565 // expected auto = const int64_t
6666 auto * x2_data = x2->data <int64_t >();
6767 // e.g get [0, 1, 2, 3, ...]
68- auto x2_lods = x2->lod ()[0 ];
68+ size_t x2_lods_size = x2->dims ()[0 ];
6969 Vector<size_t > x1_lods (1 , 0 );
7070 if (!is_x1_lod) {
7171 for (int i = 0 ; i < x1->dims ()[0 ]; i++) {
7272 x1_lods.push_back (i + 1 );
7373 }
7474 } else {
75- x1_lods = context.Input <LoDTensor>(" Ins" )->lod ()[0 ];
75+ // new: lod_level=0 => lod() return {}
76+ if (x1->lod ().size () != 0 ) {
77+ x1_lods = x1->lod ()[0 ];
78+ } else {
79+ for (int i = 0 ; i < x1->dims ()[0 ]; i++) {
80+ x1_lods.push_back (i + 1 );
81+ }
82+ }
7683 }
7784 std::unordered_map<int64_t , int64_t > mmap_aux;
7885 Vector<size_t > out_lods (1 , 0 );
79- for (size_t i = 0 ; i < x2_lods. size () - 1 ; i++) {
80- for (size_t j = x2_lods[i] ; j < x2_lods[ i + 1 ] ; j++) {
86+ for (size_t i = 0 ; i < x2_lods_size ; i++) {
87+ for (size_t j = i ; j < i + 1 ; j++) {
8188 if (filter_tag.find (x2_data[j]) != filter_tag.end ()) {
8289 size_t batch_len = x1_lods[i + 1 ] - x1_lods[i];
8390 mmap_aux[out_lods.back ()] = x1_lods[i];
0 commit comments