fix test_scale_op skipped test (#37153)

paddle/fluid/operators/scale_op.cc

@@ -73,13 +73,19 @@ class ScaleOp : public framework::OperatorWithKernel {
 
   framework::KernelSignature GetExpectedPtenKernelArgs(
       const framework::ExecutionContext &ctx) const override {
-    if (ctx.HasInput("ScaleTensor")) {
-      return framework::KernelSignature("scale.host", {"X", "ScaleTensor"},
-                                        {"bias", "bias_after_scale"}, {"Out"});
-    } else {
-      return framework::KernelSignature(
-          "scale", {"X"}, {"scale", "bias", "bias_after_scale"}, {"Out"});
+    if (ctx.InputVar("X")->IsType<framework::LoDTensor>() ||
+        ctx.InputVar("X")->IsType<framework::Tensor>()) {
+      if (ctx.HasInput("ScaleTensor")) {
+        return framework::KernelSignature("scale.host", {"X", "ScaleTensor"},
+                                          {"bias", "bias_after_scale"},
+                                          {"Out"});
+      } else {
+        return framework::KernelSignature(
+            "scale", {"X"}, {"scale", "bias", "bias_after_scale"}, {"Out"});
+      }
     }
+    // TODO(chenweihang): support other cases after selected rows added
+    return framework::KernelSignature("scale.unregistered", {}, {}, {});
   }
 };
 

python/paddle/fluid/tests/unittests/test_scale_op.py

@@ -109,9 +109,7 @@ class TestScaleOpSelectedRows(unittest.TestCase):
 
         assert (in_array * scale == result_array).all()
         assert in_height == out_height
-        # TODO(chenweihang): output rows and height cannot be shared into
-        # fluid output tensor
-        # assert in_rows == out_rows
+        assert in_rows == out_rows
 
     def test_scale_selected_rows(self):
         places = [core.CPUPlace()]