diff --git a/source/op/pt/tabulate_multi_device.cc b/source/op/pt/tabulate_multi_device.cc
index 530e9ddf4d..cede1d03d9 100644
--- a/source/op/pt/tabulate_multi_device.cc
+++ b/source/op/pt/tabulate_multi_device.cc
@@ -602,8 +602,11 @@ void TabulateFusionSeRGradGradForward(const torch::Tensor& table_tensor,
   }
 }
 
-class TabulateFusionSeAOp
-    : public torch::autograd::Function<TabulateFusionSeAOp> {
+class TabulateFusionSeAGradOp
+    : public torch::autograd::Function<TabulateFusionSeAGradOp> {
+ private:
+  std::string device;
+
  public:
   static torch::autograd::variable_list forward(
       torch::autograd::AutogradContext* ctx,
@@ -611,14 +614,16 @@ class TabulateFusionSeAOp
       const torch::Tensor& table_info_tensor,
       const torch::Tensor& em_x_tensor,
       const torch::Tensor& em_tensor,
-      int64_t last_layer_size) {
+      const torch::Tensor& dy_tensor,
+      const torch::Tensor& descriptor_tensor) {
     bool type_flag = (table_tensor.dtype() == torch::kDouble) ? true : false;
     if (type_flag) {
       return forward_t<double>(ctx, table_tensor, table_info_tensor,
-                               em_x_tensor, em_tensor, last_layer_size);
+                               em_x_tensor, em_tensor, dy_tensor,
+                               descriptor_tensor);
     } else {
       return forward_t<float>(ctx, table_tensor, table_info_tensor, em_x_tensor,
-                              em_tensor, last_layer_size);
+                              em_tensor, dy_tensor, descriptor_tensor);
     }
   }
 
@@ -629,26 +634,28 @@ class TabulateFusionSeAOp
       const torch::Tensor& table_info_tensor,
       const torch::Tensor& em_x_tensor,
       const torch::Tensor& em_tensor,
-      int64_t last_layer_size) {
-    // allocate output tensors
-    auto options = torch::TensorOptions()
-                       .dtype(table_tensor.dtype())
-                       .device(table_tensor.device());
-    torch::Tensor descriptor_tensor =
-        torch::empty({em_tensor.size(0), 4, last_layer_size}, options);
+      const torch::Tensor& dy_tensor,
+      const torch::Tensor& descriptor_tensor) {
+    // Allocate output tensors
+    torch::Tensor dy_dem_x_tensor = torch::zeros_like(em_x_tensor);
+    torch::Tensor dy_dem_tensor = torch::zeros_like(em_tensor);
+    torch::Tensor dy_dtwo_tensor = at::Tensor();
     // compute
-    TabulateFusionSeAForward<FPTYPE>(table_tensor, table_info_tensor,
-                                     em_x_tensor, em_tensor, at::Tensor(),
-                                     last_layer_size, descriptor_tensor);
+    TabulateFusionSeAGradForward<FPTYPE>(
+        table_tensor, table_info_tensor, em_x_tensor, em_tensor, at::Tensor(),
+        dy_tensor, descriptor_tensor, dy_dem_x_tensor, dy_dem_tensor,
+        dy_dtwo_tensor);
     // save data
     ctx->save_for_backward({table_tensor, table_info_tensor, em_x_tensor,
                             em_tensor, descriptor_tensor});
-    return {descriptor_tensor};
+
+    return torch::autograd::variable_list{dy_dem_x_tensor, dy_dem_tensor};
   }
 
   static torch::autograd::variable_list backward(
       torch::autograd::AutogradContext* ctx,
       torch::autograd::variable_list grad_output) {
+    // load data
     torch::autograd::variable_list saved_variables = ctx->get_saved_variables();
     torch::Tensor table_tensor = saved_variables[0];
     bool type_flag = (table_tensor.dtype() == torch::kDouble) ? true : false;
@@ -669,31 +676,32 @@ class TabulateFusionSeAOp
     torch::Tensor table_info_tensor = saved_variables[1];
     torch::Tensor em_x_tensor = saved_variables[2];
     torch::Tensor em_tensor = saved_variables[3];
-    torch::Tensor two_embed_tensor = at::Tensor();
     torch::Tensor descriptor_tensor = saved_variables[4];
 
-    // ensure the gradient output is contiguous
-    torch::Tensor dy_tensor = grad_output[0].contiguous();
+    bool is_sorted = true;
+
+    torch::Tensor dz_dy_dem_x_tensor = grad_output[0].defined()
+                                           ? grad_output[0].contiguous()
+                                           : torch::zeros_like(em_x_tensor);
+    torch::Tensor dz_dy_dem_tensor = grad_output[1].defined()
+                                         ? grad_output[1].contiguous()
+                                         : torch::zeros_like(em_tensor);
     // allocate output tensors
-    torch::Tensor dy_dem_x_tensor = torch::zeros_like(em_x_tensor);
-    torch::Tensor dy_dem_tensor = torch::zeros_like(em_tensor);
-    torch::Tensor dy_dtwo_tensor = at::Tensor();
+    torch::Tensor dz_dy_tensor = torch::empty_like(descriptor_tensor);
     // compute
-    TabulateFusionSeAGradForward<FPTYPE>(
-        table_tensor, table_info_tensor, em_x_tensor, em_tensor,
-        two_embed_tensor, dy_tensor, descriptor_tensor, dy_dem_x_tensor,
-        dy_dem_tensor, dy_dtwo_tensor);
+    TabulateFusionSeAGradGradForward<FPTYPE>(
+        table_tensor, table_info_tensor, em_x_tensor, em_tensor, at::Tensor(),
+        dz_dy_dem_x_tensor, dz_dy_dem_tensor, at::Tensor(), descriptor_tensor,
+        is_sorted, dz_dy_tensor);
 
-    return {at::Tensor(), at::Tensor(), dy_dem_x_tensor, dy_dem_tensor,
-            at::Tensor()};
+    return torch::autograd::variable_list{at::Tensor(), at::Tensor(),
+                                          at::Tensor(), at::Tensor(),
+                                          dz_dy_tensor, at::Tensor()};
   }
 };
 
-class TabulateFusionSeAGradOp
-    : public torch::autograd::Function<TabulateFusionSeAGradOp> {
- private:
-  std::string device;
-
+class TabulateFusionSeAGradGradOp
+    : public torch::autograd::Function<TabulateFusionSeAGradGradOp> {
  public:
   static torch::autograd::variable_list forward(
       torch::autograd::AutogradContext* ctx,
@@ -701,16 +709,19 @@ class TabulateFusionSeAGradOp
       const torch::Tensor& table_info_tensor,
       const torch::Tensor& em_x_tensor,
       const torch::Tensor& em_tensor,
-      const torch::Tensor& dy_tensor,
-      const torch::Tensor& descriptor_tensor) {
+      const torch::Tensor& dz_dy_dem_x_tensor,
+      const torch::Tensor& dz_dy_dem_tensor,
+      const torch::Tensor& descriptor_tensor,
+      bool is_sorted) {
     bool type_flag = (table_tensor.dtype() == torch::kDouble) ? true : false;
     if (type_flag) {
       return forward_t<double>(ctx, table_tensor, table_info_tensor,
-                               em_x_tensor, em_tensor, dy_tensor,
-                               descriptor_tensor);
+                               em_x_tensor, em_tensor, dz_dy_dem_x_tensor,
+                               dz_dy_dem_tensor, descriptor_tensor, is_sorted);
     } else {
       return forward_t<float>(ctx, table_tensor, table_info_tensor, em_x_tensor,
-                              em_tensor, dy_tensor, descriptor_tensor);
+                              em_tensor, dz_dy_dem_x_tensor, dz_dy_dem_tensor,
+                              descriptor_tensor, is_sorted);
     }
   }
 
@@ -721,28 +732,69 @@ class TabulateFusionSeAGradOp
       const torch::Tensor& table_info_tensor,
       const torch::Tensor& em_x_tensor,
       const torch::Tensor& em_tensor,
-      const torch::Tensor& dy_tensor,
-      const torch::Tensor& descriptor_tensor) {
-    // Allocate output tensors
-    torch::Tensor dy_dem_x_tensor = torch::zeros_like(em_x_tensor);
-    torch::Tensor dy_dem_tensor = torch::zeros_like(em_tensor);
-    torch::Tensor dy_dtwo_tensor = at::Tensor();
+      const torch::Tensor& dz_dy_dem_x_tensor,
+      const torch::Tensor& dz_dy_dem_tensor,
+      const torch::Tensor& descriptor_tensor,
+      bool is_sorted) {
+    // Allocate output tensor
+    torch::Tensor dz_dy_tensor = torch::empty_like(descriptor_tensor);
     // compute
-    TabulateFusionSeAGradForward<FPTYPE>(
+    TabulateFusionSeAGradGradForward<FPTYPE>(
         table_tensor, table_info_tensor, em_x_tensor, em_tensor, at::Tensor(),
-        dy_tensor, descriptor_tensor, dy_dem_x_tensor, dy_dem_tensor,
-        dy_dtwo_tensor);
+        dz_dy_dem_x_tensor, dz_dy_dem_tensor, at::Tensor(), descriptor_tensor,
+        is_sorted, dz_dy_tensor);
+
+    return torch::autograd::variable_list{dz_dy_tensor};
+  }
+};
+
+class TabulateFusionSeAOp
+    : public torch::autograd::Function<TabulateFusionSeAOp> {
+ public:
+  static torch::autograd::variable_list forward(
+      torch::autograd::AutogradContext* ctx,
+      const torch::Tensor& table_tensor,
+      const torch::Tensor& table_info_tensor,
+      const torch::Tensor& em_x_tensor,
+      const torch::Tensor& em_tensor,
+      int64_t last_layer_size) {
+    bool type_flag = (table_tensor.dtype() == torch::kDouble) ? true : false;
+    if (type_flag) {
+      return forward_t<double>(ctx, table_tensor, table_info_tensor,
+                               em_x_tensor, em_tensor, last_layer_size);
+    } else {
+      return forward_t<float>(ctx, table_tensor, table_info_tensor, em_x_tensor,
+                              em_tensor, last_layer_size);
+    }
+  }
+
+  template <typename FPTYPE>
+  static torch::autograd::variable_list forward_t(
+      torch::autograd::AutogradContext* ctx,
+      const torch::Tensor& table_tensor,
+      const torch::Tensor& table_info_tensor,
+      const torch::Tensor& em_x_tensor,
+      const torch::Tensor& em_tensor,
+      int64_t last_layer_size) {
+    // allocate output tensors
+    auto options = torch::TensorOptions()
+                       .dtype(table_tensor.dtype())
+                       .device(table_tensor.device());
+    torch::Tensor descriptor_tensor =
+        torch::empty({em_tensor.size(0), 4, last_layer_size}, options);
+    // compute
+    TabulateFusionSeAForward<FPTYPE>(table_tensor, table_info_tensor,
+                                     em_x_tensor, em_tensor, at::Tensor(),
+                                     last_layer_size, descriptor_tensor);
     // save data
     ctx->save_for_backward({table_tensor, table_info_tensor, em_x_tensor,
                             em_tensor, descriptor_tensor});
-
-    return torch::autograd::variable_list{dy_dem_x_tensor, dy_dem_tensor};
+    return {descriptor_tensor};
   }
 
   static torch::autograd::variable_list backward(
       torch::autograd::AutogradContext* ctx,
       torch::autograd::variable_list grad_output) {
-    // load data
     torch::autograd::variable_list saved_variables = ctx->get_saved_variables();
     torch::Tensor table_tensor = saved_variables[0];
     bool type_flag = (table_tensor.dtype() == torch::kDouble) ? true : false;
@@ -765,26 +817,20 @@ class TabulateFusionSeAGradOp
     torch::Tensor em_tensor = saved_variables[3];
     torch::Tensor descriptor_tensor = saved_variables[4];
 
-    bool is_sorted = true;
-
-    torch::Tensor dz_dy_dem_x_tensor = grad_output[0].contiguous();
-    torch::Tensor dz_dy_dem_tensor = grad_output[1].contiguous();
-    // allocate output tensors
-    torch::Tensor dz_dy_tensor = torch::empty_like(descriptor_tensor);
-    // compute
-    TabulateFusionSeAGradGradForward<FPTYPE>(
-        table_tensor, table_info_tensor, em_x_tensor, em_tensor, at::Tensor(),
-        dz_dy_dem_x_tensor, dz_dy_dem_tensor, at::Tensor(), descriptor_tensor,
-        is_sorted, dz_dy_tensor);
+    // ensure the gradient output is contiguous
+    torch::Tensor dy_tensor = grad_output[0].contiguous();
+    torch::autograd::variable_list dy_dem_tensors =
+        TabulateFusionSeAGradOp::apply(table_tensor, table_info_tensor,
+                                       em_x_tensor, em_tensor, dy_tensor,
+                                       descriptor_tensor);
 
-    return torch::autograd::variable_list{at::Tensor(), at::Tensor(),
-                                          at::Tensor(), at::Tensor(),
-                                          dz_dy_tensor, at::Tensor()};
+    return {at::Tensor(), at::Tensor(), dy_dem_tensors[0], dy_dem_tensors[1],
+            at::Tensor()};
   }
 };
 
-class TabulateFusionSeAGradGradOp
-    : public torch::autograd::Function<TabulateFusionSeAGradGradOp> {
+class TabulateFusionSeAttenGradOp
+    : public torch::autograd::Function<TabulateFusionSeAttenGradOp> {
  public:
   static torch::autograd::variable_list forward(
       torch::autograd::AutogradContext* ctx,
@@ -792,18 +838,18 @@ class TabulateFusionSeAGradGradOp
       const torch::Tensor& table_info_tensor,
       const torch::Tensor& em_x_tensor,
       const torch::Tensor& em_tensor,
-      const torch::Tensor& dz_dy_dem_x_tensor,
-      const torch::Tensor& dz_dy_dem_tensor,
+      const torch::Tensor& two_embed_tensor,
+      const torch::Tensor& dy_tensor,
       const torch::Tensor& descriptor_tensor,
       bool is_sorted) {
     bool type_flag = (table_tensor.dtype() == torch::kDouble) ? true : false;
     if (type_flag) {
       return forward_t<double>(ctx, table_tensor, table_info_tensor,
-                               em_x_tensor, em_tensor, dz_dy_dem_x_tensor,
-                               dz_dy_dem_tensor, descriptor_tensor, is_sorted);
+                               em_x_tensor, em_tensor, two_embed_tensor,
+                               dy_tensor, descriptor_tensor, is_sorted);
     } else {
       return forward_t<float>(ctx, table_tensor, table_info_tensor, em_x_tensor,
-                              em_tensor, dz_dy_dem_x_tensor, dz_dy_dem_tensor,
+                              em_tensor, two_embed_tensor, dy_tensor,
                               descriptor_tensor, is_sorted);
     }
   }
@@ -815,19 +861,70 @@ class TabulateFusionSeAGradGradOp
       const torch::Tensor& table_info_tensor,
       const torch::Tensor& em_x_tensor,
       const torch::Tensor& em_tensor,
-      const torch::Tensor& dz_dy_dem_x_tensor,
-      const torch::Tensor& dz_dy_dem_tensor,
+      const torch::Tensor& two_embed_tensor,
+      const torch::Tensor& dy_tensor,
       const torch::Tensor& descriptor_tensor,
       bool is_sorted) {
-    // Allocate output tensor
+    torch::Tensor dy_dem_x_tensor = torch::zeros_like(em_x_tensor);
+    torch::Tensor dy_dem_tensor = torch::zeros_like(em_tensor);
+    torch::Tensor dy_dtwo_tensor = torch::zeros_like(two_embed_tensor);
+    TabulateFusionSeAGradForward<FPTYPE>(
+        table_tensor, table_info_tensor, em_x_tensor, em_tensor,
+        two_embed_tensor, dy_tensor, descriptor_tensor, dy_dem_x_tensor,
+        dy_dem_tensor, dy_dtwo_tensor);
+
+    ctx->save_for_backward({table_tensor, table_info_tensor, em_x_tensor,
+                            em_tensor, two_embed_tensor, descriptor_tensor});
+    ctx->saved_data["is_sorted"] = is_sorted;
+
+    return torch::autograd::variable_list{dy_dem_x_tensor, dy_dem_tensor,
+                                          dy_dtwo_tensor};
+  }
+
+  static torch::autograd::variable_list backward(
+      torch::autograd::AutogradContext* ctx,
+      torch::autograd::variable_list grad_output) {
+    torch::autograd::variable_list saved_variables = ctx->get_saved_variables();
+    torch::Tensor table_tensor = saved_variables[0];
+    bool type_flag = (table_tensor.dtype() == torch::kDouble) ? true : false;
+    if (type_flag) {
+      return backward_t<double>(ctx, grad_output);
+    } else {
+      return backward_t<float>(ctx, grad_output);
+    }
+  }
+
+  template <typename FPTYPE>
+  static torch::autograd::variable_list backward_t(
+      torch::autograd::AutogradContext* ctx,
+      torch::autograd::variable_list grad_output) {
+    torch::autograd::variable_list saved_variables = ctx->get_saved_variables();
+    torch::Tensor table_tensor = saved_variables[0];
+    torch::Tensor table_info_tensor = saved_variables[1];
+    torch::Tensor em_x_tensor = saved_variables[2];
+    torch::Tensor em_tensor = saved_variables[3];
+    torch::Tensor two_embed_tensor = saved_variables[4];
+    torch::Tensor descriptor_tensor = saved_variables[5];
+    bool is_sorted = ctx->saved_data["is_sorted"].toBool();
+
+    torch::Tensor dz_dy_dem_x_tensor = grad_output[0].defined()
+                                           ? grad_output[0].contiguous()
+                                           : torch::zeros_like(em_x_tensor);
+    torch::Tensor dz_dy_dem_tensor = grad_output[1].defined()
+                                         ? grad_output[1].contiguous()
+                                         : torch::zeros_like(em_tensor);
+    torch::Tensor dz_dy_dtwo_tensor = grad_output[2].defined()
+                                          ? grad_output[2].contiguous()
+                                          : torch::zeros_like(two_embed_tensor);
     torch::Tensor dz_dy_tensor = torch::empty_like(descriptor_tensor);
-    // compute
     TabulateFusionSeAGradGradForward<FPTYPE>(
-        table_tensor, table_info_tensor, em_x_tensor, em_tensor, at::Tensor(),
-        dz_dy_dem_x_tensor, dz_dy_dem_tensor, at::Tensor(), descriptor_tensor,
-        is_sorted, dz_dy_tensor);
+        table_tensor, table_info_tensor, em_x_tensor, em_tensor,
+        two_embed_tensor, dz_dy_dem_x_tensor, dz_dy_dem_tensor,
+        dz_dy_dtwo_tensor, descriptor_tensor, is_sorted, dz_dy_tensor);
 
-    return torch::autograd::variable_list{dz_dy_tensor};
+    return torch::autograd::variable_list{
+        at::Tensor(), at::Tensor(), at::Tensor(), at::Tensor(),
+        at::Tensor(), dz_dy_tensor, at::Tensor(), at::Tensor()};
   }
 };
 
@@ -878,6 +975,7 @@ class TabulateFusionSeAttenOp
     // save data
     ctx->save_for_backward({table_tensor, table_info_tensor, em_x_tensor,
                             em_tensor, two_embed_tensor, descriptor_tensor});
+    ctx->saved_data["is_sorted"] = is_sorted;
     return {descriptor_tensor};
   }
 
@@ -906,20 +1004,101 @@ class TabulateFusionSeAttenOp
     torch::Tensor em_tensor = saved_variables[3];
     torch::Tensor two_embed_tensor = saved_variables[4];
     torch::Tensor descriptor_tensor = saved_variables[5];
+    bool is_sorted = ctx->saved_data["is_sorted"].toBool();
 
     torch::Tensor dy_tensor = grad_output[0].contiguous();
-    // allocate output tensors
+    torch::autograd::variable_list dy_dem_tensors =
+        TabulateFusionSeAttenGradOp::apply(
+            table_tensor, table_info_tensor, em_x_tensor, em_tensor,
+            two_embed_tensor, dy_tensor, descriptor_tensor, is_sorted);
+
+    return {at::Tensor(),      at::Tensor(),      dy_dem_tensors[0],
+            dy_dem_tensors[1], dy_dem_tensors[2], at::Tensor(),
+            at::Tensor()};
+  }
+};
+
+class TabulateFusionSeTGradOp
+    : public torch::autograd::Function<TabulateFusionSeTGradOp> {
+ public:
+  static torch::autograd::variable_list forward(
+      torch::autograd::AutogradContext* ctx,
+      const torch::Tensor& table_tensor,
+      const torch::Tensor& table_info_tensor,
+      const torch::Tensor& em_x_tensor,
+      const torch::Tensor& em_tensor,
+      const torch::Tensor& dy_tensor,
+      const torch::Tensor& descriptor_tensor) {
+    bool type_flag = (table_tensor.dtype() == torch::kDouble) ? true : false;
+    if (type_flag) {
+      return forward_t<double>(ctx, table_tensor, table_info_tensor,
+                               em_x_tensor, em_tensor, dy_tensor,
+                               descriptor_tensor);
+    } else {
+      return forward_t<float>(ctx, table_tensor, table_info_tensor, em_x_tensor,
+                              em_tensor, dy_tensor, descriptor_tensor);
+    }
+  }
+
+  template <typename FPTYPE>
+  static torch::autograd::variable_list forward_t(
+      torch::autograd::AutogradContext* ctx,
+      const torch::Tensor& table_tensor,
+      const torch::Tensor& table_info_tensor,
+      const torch::Tensor& em_x_tensor,
+      const torch::Tensor& em_tensor,
+      const torch::Tensor& dy_tensor,
+      const torch::Tensor& descriptor_tensor) {
     torch::Tensor dy_dem_x_tensor = torch::zeros_like(em_x_tensor);
     torch::Tensor dy_dem_tensor = torch::zeros_like(em_tensor);
-    torch::Tensor dy_dtwo_tensor = torch::zeros_like(two_embed_tensor);
-    // compute
-    TabulateFusionSeAGradForward<FPTYPE>(
+    TabulateFusionSeTGradForward<FPTYPE>(
+        table_tensor, table_info_tensor, em_x_tensor, em_tensor, dy_tensor,
+        descriptor_tensor, dy_dem_x_tensor, dy_dem_tensor);
+
+    ctx->save_for_backward({table_tensor, table_info_tensor, em_x_tensor,
+                            em_tensor, descriptor_tensor});
+
+    return torch::autograd::variable_list{dy_dem_x_tensor, dy_dem_tensor};
+  }
+
+  static torch::autograd::variable_list backward(
+      torch::autograd::AutogradContext* ctx,
+      torch::autograd::variable_list grad_output) {
+    torch::autograd::variable_list saved_variables = ctx->get_saved_variables();
+    torch::Tensor table_tensor = saved_variables[0];
+    bool type_flag = (table_tensor.dtype() == torch::kDouble) ? true : false;
+    if (type_flag) {
+      return backward_t<double>(ctx, grad_output);
+    } else {
+      return backward_t<float>(ctx, grad_output);
+    }
+  }
+
+  template <typename FPTYPE>
+  static torch::autograd::variable_list backward_t(
+      torch::autograd::AutogradContext* ctx,
+      torch::autograd::variable_list grad_output) {
+    torch::autograd::variable_list saved_variables = ctx->get_saved_variables();
+    torch::Tensor table_tensor = saved_variables[0];
+    torch::Tensor table_info_tensor = saved_variables[1];
+    torch::Tensor em_x_tensor = saved_variables[2];
+    torch::Tensor em_tensor = saved_variables[3];
+    torch::Tensor descriptor_tensor = saved_variables[4];
+
+    torch::Tensor dz_dy_dem_x_tensor = grad_output[0].defined()
+                                           ? grad_output[0].contiguous()
+                                           : torch::zeros_like(em_x_tensor);
+    torch::Tensor dz_dy_dem_tensor = grad_output[1].defined()
+                                         ? grad_output[1].contiguous()
+                                         : torch::zeros_like(em_tensor);
+    torch::Tensor dz_dy_tensor = torch::empty_like(descriptor_tensor);
+    TabulateFusionSeTGradGradForward<FPTYPE>(
         table_tensor, table_info_tensor, em_x_tensor, em_tensor,
-        two_embed_tensor, dy_tensor, descriptor_tensor, dy_dem_x_tensor,
-        dy_dem_tensor, dy_dtwo_tensor);
+        dz_dy_dem_x_tensor, dz_dy_dem_tensor, descriptor_tensor, dz_dy_tensor);
 
-    return {at::Tensor(),   at::Tensor(), dy_dem_x_tensor, dy_dem_tensor,
-            dy_dtwo_tensor, at::Tensor(), at::Tensor()};
+    return torch::autograd::variable_list{at::Tensor(), at::Tensor(),
+                                          at::Tensor(), at::Tensor(),
+                                          dz_dy_tensor, at::Tensor()};
   }
 };
 
@@ -993,19 +1172,91 @@ class TabulateFusionSeTOp
     torch::Tensor descriptor_tensor = saved_variables[4];
 
     torch::Tensor dy_tensor = grad_output[0].contiguous();
-    // allocate output tensors
-    torch::Tensor dy_dem_x_tensor = torch::zeros_like(em_x_tensor);
-    torch::Tensor dy_dem_tensor = torch::zeros_like(em_tensor);
-    // compute
-    TabulateFusionSeTGradForward<FPTYPE>(
-        table_tensor, table_info_tensor, em_x_tensor, em_tensor, dy_tensor,
-        descriptor_tensor, dy_dem_x_tensor, dy_dem_tensor);
+    torch::autograd::variable_list dy_dem_tensors =
+        TabulateFusionSeTGradOp::apply(table_tensor, table_info_tensor,
+                                       em_x_tensor, em_tensor, dy_tensor,
+                                       descriptor_tensor);
 
-    return {at::Tensor(), at::Tensor(), dy_dem_x_tensor, dy_dem_tensor,
+    return {at::Tensor(), at::Tensor(), dy_dem_tensors[0], dy_dem_tensors[1],
             at::Tensor()};
   }
 };
 
+class TabulateFusionSeRGradOp
+    : public torch::autograd::Function<TabulateFusionSeRGradOp> {
+ public:
+  static torch::autograd::variable_list forward(
+      torch::autograd::AutogradContext* ctx,
+      const torch::Tensor& table_tensor,
+      const torch::Tensor& table_info_tensor,
+      const torch::Tensor& em_tensor,
+      const torch::Tensor& dy_tensor,
+      const torch::Tensor& descriptor_tensor) {
+    bool type_flag = (table_tensor.dtype() == torch::kDouble) ? true : false;
+    if (type_flag) {
+      return forward_t<double>(ctx, table_tensor, table_info_tensor, em_tensor,
+                               dy_tensor, descriptor_tensor);
+    } else {
+      return forward_t<float>(ctx, table_tensor, table_info_tensor, em_tensor,
+                              dy_tensor, descriptor_tensor);
+    }
+  }
+
+  template <typename FPTYPE>
+  static torch::autograd::variable_list forward_t(
+      torch::autograd::AutogradContext* ctx,
+      const torch::Tensor& table_tensor,
+      const torch::Tensor& table_info_tensor,
+      const torch::Tensor& em_tensor,
+      const torch::Tensor& dy_tensor,
+      const torch::Tensor& descriptor_tensor) {
+    torch::Tensor dy_dem_tensor = torch::zeros_like(em_tensor);
+    TabulateFusionSeRGradForward<FPTYPE>(table_tensor, table_info_tensor,
+                                         em_tensor, dy_tensor,
+                                         descriptor_tensor, dy_dem_tensor);
+
+    ctx->save_for_backward(
+        {table_tensor, table_info_tensor, em_tensor, descriptor_tensor});
+
+    return torch::autograd::variable_list{dy_dem_tensor};
+  }
+
+  static torch::autograd::variable_list backward(
+      torch::autograd::AutogradContext* ctx,
+      torch::autograd::variable_list grad_output) {
+    torch::autograd::variable_list saved_variables = ctx->get_saved_variables();
+    torch::Tensor table_tensor = saved_variables[0];
+    bool type_flag = (table_tensor.dtype() == torch::kDouble) ? true : false;
+    if (type_flag) {
+      return backward_t<double>(ctx, grad_output);
+    } else {
+      return backward_t<float>(ctx, grad_output);
+    }
+  }
+
+  template <typename FPTYPE>
+  static torch::autograd::variable_list backward_t(
+      torch::autograd::AutogradContext* ctx,
+      torch::autograd::variable_list grad_output) {
+    torch::autograd::variable_list saved_variables = ctx->get_saved_variables();
+    torch::Tensor table_tensor = saved_variables[0];
+    torch::Tensor table_info_tensor = saved_variables[1];
+    torch::Tensor em_tensor = saved_variables[2];
+    torch::Tensor descriptor_tensor = saved_variables[3];
+
+    torch::Tensor dz_dy_dem_tensor = grad_output[0].defined()
+                                         ? grad_output[0].contiguous()
+                                         : torch::zeros_like(em_tensor);
+    torch::Tensor dz_dy_tensor = torch::empty_like(descriptor_tensor);
+    TabulateFusionSeRGradGradForward<FPTYPE>(table_tensor, table_info_tensor,
+                                             em_tensor, dz_dy_dem_tensor,
+                                             descriptor_tensor, dz_dy_tensor);
+
+    return torch::autograd::variable_list{
+        at::Tensor(), at::Tensor(), at::Tensor(), dz_dy_tensor, at::Tensor()};
+  }
+};
+
 class TabulateFusionSeROp
     : public torch::autograd::Function<TabulateFusionSeROp> {
  public:
@@ -1072,14 +1323,93 @@ class TabulateFusionSeROp
     torch::Tensor descriptor_tensor = saved_variables[3];
 
     torch::Tensor dy_tensor = grad_output[0].contiguous();
-    // allocate output tensors
-    torch::Tensor dy_dem_tensor = torch::zeros_like(em_tensor);
-    // compute
-    TabulateFusionSeRGradForward<FPTYPE>(table_tensor, table_info_tensor,
-                                         em_tensor, dy_tensor,
-                                         descriptor_tensor, dy_dem_tensor);
+    torch::autograd::variable_list dy_dem_tensors =
+        TabulateFusionSeRGradOp::apply(table_tensor, table_info_tensor,
+                                       em_tensor, dy_tensor, descriptor_tensor);
+
+    return {at::Tensor(), at::Tensor(), dy_dem_tensors[0], at::Tensor()};
+  }
+};
+
+class TabulateFusionSeTTebdGradOp
+    : public torch::autograd::Function<TabulateFusionSeTTebdGradOp> {
+ public:
+  static torch::autograd::variable_list forward(
+      torch::autograd::AutogradContext* ctx,
+      const torch::Tensor& table_tensor,
+      const torch::Tensor& table_info_tensor,
+      const torch::Tensor& em_x_tensor,
+      const torch::Tensor& em_tensor,
+      const torch::Tensor& dy_tensor,
+      const torch::Tensor& descriptor_tensor) {
+    bool type_flag = (table_tensor.dtype() == torch::kDouble) ? true : false;
+    if (type_flag) {
+      return forward_t<double>(ctx, table_tensor, table_info_tensor,
+                               em_x_tensor, em_tensor, dy_tensor,
+                               descriptor_tensor);
+    } else {
+      return forward_t<float>(ctx, table_tensor, table_info_tensor, em_x_tensor,
+                              em_tensor, dy_tensor, descriptor_tensor);
+    }
+  }
+
+  template <typename FPTYPE>
+  static torch::autograd::variable_list forward_t(
+      torch::autograd::AutogradContext* ctx,
+      const torch::Tensor& table_tensor,
+      const torch::Tensor& table_info_tensor,
+      const torch::Tensor& em_x_tensor,
+      const torch::Tensor& em_tensor,
+      const torch::Tensor& dy_tensor,
+      const torch::Tensor& descriptor_tensor) {
+    torch::Tensor dy_dem_x_tensor = torch::zeros_like(em_x_tensor);
+    TabulateFusionSeTTebdGradForward<FPTYPE>(
+        table_tensor, table_info_tensor, em_x_tensor, em_tensor, dy_tensor,
+        descriptor_tensor, dy_dem_x_tensor);
+
+    ctx->save_for_backward({table_tensor, table_info_tensor, em_x_tensor,
+                            em_tensor, descriptor_tensor});
+
+    return torch::autograd::variable_list{dy_dem_x_tensor};
+  }
+
+  static torch::autograd::variable_list backward(
+      torch::autograd::AutogradContext* ctx,
+      torch::autograd::variable_list grad_output) {
+    torch::autograd::variable_list saved_variables = ctx->get_saved_variables();
+    torch::Tensor table_tensor = saved_variables[0];
+    bool type_flag = (table_tensor.dtype() == torch::kDouble) ? true : false;
+    if (type_flag) {
+      return backward_t<double>(ctx, grad_output);
+    } else {
+      return backward_t<float>(ctx, grad_output);
+    }
+  }
+
+  template <typename FPTYPE>
+  static torch::autograd::variable_list backward_t(
+      torch::autograd::AutogradContext* ctx,
+      torch::autograd::variable_list grad_output) {
+    torch::autograd::variable_list saved_variables = ctx->get_saved_variables();
+    torch::Tensor table_tensor = saved_variables[0];
+    torch::Tensor table_info_tensor = saved_variables[1];
+    torch::Tensor em_x_tensor = saved_variables[2];
+    torch::Tensor em_tensor = saved_variables[3];
+    torch::Tensor descriptor_tensor = saved_variables[4];
 
-    return {at::Tensor(), at::Tensor(), dy_dem_tensor, at::Tensor()};
+    torch::Tensor dz_dy_dem_x_tensor =
+        grad_output[0].defined()
+            ? grad_output[0].contiguous().view(
+                  {em_tensor.size(0), em_tensor.size(1), em_tensor.size(2)})
+            : torch::zeros_like(em_tensor);
+    torch::Tensor dz_dy_tensor = torch::empty_like(descriptor_tensor);
+    TabulateFusionSeTTebdGradGradForward<FPTYPE>(
+        table_tensor, table_info_tensor, em_x_tensor, em_tensor,
+        dz_dy_dem_x_tensor, descriptor_tensor, dz_dy_tensor);
+
+    return torch::autograd::variable_list{at::Tensor(), at::Tensor(),
+                                          at::Tensor(), at::Tensor(),
+                                          dz_dy_tensor, at::Tensor()};
   }
 };
 
@@ -1155,14 +1485,12 @@ class TabulateFusionSeTTebdOp
     torch::Tensor descriptor_tensor = saved_variables[4];
 
     torch::Tensor dy_tensor = grad_output[0].contiguous();
-    // allocate output tensors
-    torch::Tensor dy_dem_x_tensor = torch::zeros_like(em_x_tensor);
-    // compute
-    TabulateFusionSeTTebdGradForward<FPTYPE>(
-        table_tensor, table_info_tensor, em_x_tensor, em_tensor, dy_tensor,
-        descriptor_tensor, dy_dem_x_tensor);
+    torch::autograd::variable_list dy_dem_tensors =
+        TabulateFusionSeTTebdGradOp::apply(table_tensor, table_info_tensor,
+                                           em_x_tensor, em_tensor, dy_tensor,
+                                           descriptor_tensor);
 
-    return {at::Tensor(), at::Tensor(), dy_dem_x_tensor, at::Tensor(),
+    return {at::Tensor(), at::Tensor(), dy_dem_tensors[0], at::Tensor(),
             at::Tensor()};
   }
 };
diff --git a/source/tests/pt/tabulate_test_utils.py b/source/tests/pt/tabulate_test_utils.py
new file mode 100644
index 0000000000..07eea61a0d
--- /dev/null
+++ b/source/tests/pt/tabulate_test_utils.py
@@ -0,0 +1,92 @@
+# SPDX-License-Identifier: LGPL-3.0-or-later
+from collections.abc import (
+    Sequence,
+)
+
+import torch
+
+
+def _nonuniform_like(tensor: torch.Tensor) -> torch.Tensor:
+    """Create deterministic nonuniform weights for gradient projections."""
+    return torch.linspace(
+        -0.5,
+        0.75,
+        tensor.numel(),
+        dtype=tensor.dtype,
+        device=tensor.device,
+    ).reshape_as(tensor)
+
+
+def _project_first_grads(
+    descriptor_tensor: torch.Tensor,
+    grad_inputs: Sequence[torch.Tensor],
+    dy_tensor: torch.Tensor,
+    weights: Sequence[torch.Tensor],
+) -> torch.Tensor:
+    first_grads = torch.autograd.grad(
+        descriptor_tensor,
+        grad_inputs,
+        grad_outputs=dy_tensor,
+        retain_graph=True,
+    )
+    projection = descriptor_tensor.new_zeros(())
+    for grad, weight in zip(first_grads, weights, strict=True):
+        projection = projection + (grad * weight).sum()
+    return projection
+
+
+def assert_second_order_backward_matches_finite_difference(
+    descriptor_tensor: torch.Tensor,
+    grad_inputs: torch.Tensor | Sequence[torch.Tensor],
+) -> None:
+    """Compare dz/dy from the grad-grad kernel with a finite-difference reference."""
+    if isinstance(grad_inputs, torch.Tensor):
+        grad_inputs = (grad_inputs,)
+
+    dy_tensor = _nonuniform_like(descriptor_tensor).requires_grad_(True)
+    first_grads = torch.autograd.grad(
+        descriptor_tensor,
+        grad_inputs,
+        grad_outputs=dy_tensor,
+        create_graph=True,
+        retain_graph=True,
+    )
+    weights = tuple(_nonuniform_like(grad) for grad in first_grads)
+    projection = descriptor_tensor.new_zeros(())
+    for grad, weight in zip(first_grads, weights, strict=True):
+        projection = projection + (grad * weight).sum()
+
+    dz_dy_tensor = torch.autograd.grad(
+        projection,
+        dy_tensor,
+        retain_graph=True,
+    )[0]
+
+    eps, atol, rtol = (
+        (1e-6, 1e-6, 1e-6)
+        if descriptor_tensor.dtype == torch.float64
+        else (1e-2, 5e-3, 5e-3)
+    )
+    reference = torch.empty_like(dy_tensor)
+    base_dy = dy_tensor.detach()
+    reference_flat = reference.view(-1)
+    for idx in range(dy_tensor.numel()):
+        dy_plus = base_dy.clone()
+        dy_minus = base_dy.clone()
+        dy_plus.view(-1)[idx] += eps
+        dy_minus.view(-1)[idx] -= eps
+
+        plus_projection = _project_first_grads(
+            descriptor_tensor, grad_inputs, dy_plus, weights
+        )
+        minus_projection = _project_first_grads(
+            descriptor_tensor, grad_inputs, dy_minus, weights
+        )
+        reference_flat[idx] = (plus_projection - minus_projection) / (2.0 * eps)
+
+    torch.testing.assert_close(
+        dz_dy_tensor,
+        reference,
+        atol=atol,
+        rtol=rtol,
+    )
diff --git a/source/tests/pt/test_tabulate_fusion_se_a.py b/source/tests/pt/test_tabulate_fusion_se_a.py
index 8861f0564e..690f25a587 100644
--- a/source/tests/pt/test_tabulate_fusion_se_a.py
+++ b/source/tests/pt/test_tabulate_fusion_se_a.py
@@ -13,6 +13,9 @@
 from ..consistent.common import (
     parameterized,
 )
+from .tabulate_test_utils import (
+    assert_second_order_backward_matches_finite_difference,
+)
 
 
 @parameterized((torch.float64, torch.float32))
@@ -1488,6 +1491,21 @@ def test_backward(self) -> None:
             rtol=self.prec,
         )
 
+    def test_second_order_backward(self) -> None:
+        forward_result = torch.ops.deepmd.tabulate_fusion_se_a(
+            self.table_tensor,
+            self.table_info_tensor,
+            self.em_x_tensor,
+            self.em_tensor,
+            self.last_layer_size,
+        )
+
+        descriptor_tensor = forward_result[0]
+        assert_second_order_backward_matches_finite_difference(
+            descriptor_tensor,
+            (self.em_x_tensor, self.em_tensor),
+        )
+
 
 if __name__ == "__main__":
     unittest.main()
diff --git a/source/tests/pt/test_tabulate_fusion_se_atten.py b/source/tests/pt/test_tabulate_fusion_se_atten.py
index 09b06dcc89..f46d0b1761 100644
--- a/source/tests/pt/test_tabulate_fusion_se_atten.py
+++ b/source/tests/pt/test_tabulate_fusion_se_atten.py
@@ -13,6 +13,9 @@
 from ..consistent.common import (
     parameterized,
 )
+from .tabulate_test_utils import (
+    assert_second_order_backward_matches_finite_difference,
+)
 
 
 @parameterized((torch.float64, torch.float32))
@@ -1624,6 +1627,23 @@ def test_backward(self) -> None:
             self.em_tensor.grad, self.expected_dy_dem, atol=self.prec, rtol=self.prec
         )
 
+    def test_second_order_backward(self) -> None:
+        forward_result = torch.ops.deepmd.tabulate_fusion_se_atten(
+            self.table_tensor,
+            self.table_info_tensor,
+            self.em_x_tensor,
+            self.em_tensor,
+            self.two_embed_tensor,
+            self.last_layer_size,
+            self.is_sorted,
+        )
+
+        descriptor_tensor = forward_result[0]
+        assert_second_order_backward_matches_finite_difference(
+            descriptor_tensor,
+            (self.em_x_tensor, self.em_tensor, self.two_embed_tensor),
+        )
+
 
 if __name__ == "__main__":
     unittest.main()
diff --git a/source/tests/pt/test_tabulate_fusion_se_r.py b/source/tests/pt/test_tabulate_fusion_se_r.py
index 34d8dbf182..d763e1fe60 100644
--- a/source/tests/pt/test_tabulate_fusion_se_r.py
+++ b/source/tests/pt/test_tabulate_fusion_se_r.py
@@ -13,6 +13,9 @@
 from ..consistent.common import (
     parameterized,
 )
+from .tabulate_test_utils import (
+    assert_second_order_backward_matches_finite_difference,
+)
 
 
 @parameterized((torch.float64, torch.float32))
@@ -1332,6 +1335,20 @@ def test_backward(self) -> None:
             self.em_tensor.grad, self.expected_dy_dem, atol=self.prec, rtol=self.prec
         )
 
+    def test_second_order_backward(self) -> None:
+        forward_result = torch.ops.deepmd.tabulate_fusion_se_r(
+            self.table_tensor,
+            self.table_info_tensor,
+            self.em_tensor,
+            self.last_layer_size,
+        )
+
+        descriptor_tensor = forward_result[0]
+        assert_second_order_backward_matches_finite_difference(
+            descriptor_tensor,
+            self.em_tensor,
+        )
+
 
 if __name__ == "__main__":
     unittest.main()
diff --git a/source/tests/pt/test_tabulate_fusion_se_t.py b/source/tests/pt/test_tabulate_fusion_se_t.py
index 50654d557d..421f7343ff 100644
--- a/source/tests/pt/test_tabulate_fusion_se_t.py
+++ b/source/tests/pt/test_tabulate_fusion_se_t.py
@@ -13,6 +13,9 @@
 from ..consistent.common import (
     parameterized,
 )
+from .tabulate_test_utils import (
+    assert_second_order_backward_matches_finite_difference,
+)
 
 
 @parameterized((torch.float64, torch.float32))
@@ -1748,6 +1751,21 @@ def test_backward(self) -> None:
             rtol=self.prec,
         )
 
+    def test_second_order_backward(self) -> None:
+        forward_result = torch.ops.deepmd.tabulate_fusion_se_t(
+            self.table_tensor,
+            self.table_info_tensor,
+            self.em_x_tensor,
+            self.em_tensor,
+            self.last_layer_size,
+        )
+
+        descriptor_tensor = forward_result[0]
+        assert_second_order_backward_matches_finite_difference(
+            descriptor_tensor,
+            (self.em_x_tensor, self.em_tensor),
+        )
+
 
 if __name__ == "__main__":
     unittest.main()
diff --git a/source/tests/pt/test_tabulate_fusion_se_t_tebd.py b/source/tests/pt/test_tabulate_fusion_se_t_tebd.py
index 10bf48c46f..49563e4adc 100644
--- a/source/tests/pt/test_tabulate_fusion_se_t_tebd.py
+++ b/source/tests/pt/test_tabulate_fusion_se_t_tebd.py
@@ -13,6 +13,9 @@
 from ..consistent.common import (
     parameterized,
 )
+from .tabulate_test_utils import (
+    assert_second_order_backward_matches_finite_difference,
+)
 
 
 @parameterized((torch.float64, torch.float32))
@@ -1913,6 +1916,21 @@ def test_backward(self) -> None:
             rtol=self.prec,
         )
 
+    def test_second_order_backward(self) -> None:
+        forward_result = torch.ops.deepmd.tabulate_fusion_se_t_tebd(
+            self.table_tensor,
+            self.table_info_tensor,
+            self.em_x_tensor,
+            self.em_tensor,
+            self.last_layer_size,
+        )
+
+        descriptor_tensor = forward_result[0]
+        assert_second_order_backward_matches_finite_difference(
+            descriptor_tensor,
+            self.em_x_tensor,
+        )
+
 
 if __name__ == "__main__":
     unittest.main()