     const Tensor* x3 = context.Input<Tensor>("Filter_tag");
     const int64_t* x3_data = x3->data<int64_t>();
 
-    // int x2_lods_filled = 1;
-
     Vector<size_t> x2_lods;
-    // Vector, in GPU
     if (x2->lod().size() != 0) {  // lod_level = 1
       x2_lods = x2->lod()[0];
-      // x2_lods_filled = 1;
-
     } else {  // lod_level = 0
       const size_t x2_lods_size = x2->dims()[0];
+      const size_t instag_per_num = x2->dims()[1];
       // x2_lods.resize(x2->dims()[0] + 1);
       // move to cuda
       x2_lods.push_back(0);
       for (size_t i = 0; i < x2_lods_size; i++) {
-        x2_lods.push_back(i + 1);
+        x2_lods.push_back(x2_lods.back() + instag_per_num);
       }
     }
 

    LoDTensor* loss_weight = context.Output<LoDTensor>("LossWeight");

    int out_first = x1_lods.back();
    // int out_first = x1->dims()[0];
    // if (x1_lods_filled) {
    //  out_first = x1_lods.back();
    // }

    out->Resize(phi::make_ddim({(int64_t)out_first, (int64_t)x1_embed_size}));
    map->Resize(phi::make_ddim({(int64_t)x2_lods_size, 3}));

    if (sum_out_lods4 > 1) {
      int out_data_num = sum_out_lods4 - 1;
      int out_start = ins_start;
      if (out_start < out_data_num) {
        int out_end = ins_end >= out_data_num ? out_data_num : ins_end;
        for (int p = out_start; p < out_end; p++) {
