Unbreak optimized sub in the case where one input is a scalar and dtype is mixed or both are Half (#5894)

Summary:
Pull Request resolved: #5894

I misplaced a return. Caught this during development in div, but not sub.
ghstack-source-id: 246398371
exported-using-ghexport

Reviewed By: kirklandsign

Differential Revision: D63919553

fbshipit-source-id: 84e27efc400711ae0253e48a4dbdb4b419140925

Author: swolchok

kernels/optimized/cpu/op_sub.cpp

@@ -134,8 +134,8 @@ Tensor& opt_sub_out(
           }
         });
       });
+      return out;
     }
-    return out;
   }
 
   auto selected_optimized_path = select_optimized_path(a, b, out);

kernels/test/op_sub_test.cpp

@@ -107,6 +107,27 @@ class OpSubOutTest : public OperatorTest {
 
 #undef ENUMERATE_TEST_ENTRY
   }
+
+  template <ScalarType DTYPE>
+  void test_broadcast_rank1_scalar() {
+    TensorFactory<DTYPE> tf;
+
+    Tensor a = tf.make({2, 1, 3}, {2, 3, 4, 5, 6, 7});
+    Tensor b = tf.make({1}, {2});
+
+    // Destination for the broadcasting div. Follow the broadcasting rules in
+    // https://fburl.com/n9wl4d0o
+    Tensor out = tf.zeros({2, 1, 3});
+
+    op_sub_out(a, b, 1, out);
+
+    Tensor ret = tf.make({2, 1, 3}, {0, 1, 2, 3, 4, 5});
+    EXPECT_TENSOR_EQ(out, ret);
+
+    op_sub_out(b, a, 1, out);
+    ret = tf.make({2, 1, 3}, {0, -1, -2, -3, -4, -5});
+    EXPECT_TENSOR_EQ(out, ret);
+  }
 };
 
 class OpSubScalarOutTest : public OperatorTest {
@@ -171,19 +192,8 @@ TEST_F(OpSubOutTest, BroadcastSupported2) {
 }
 
 TEST_F(OpSubOutTest, BroadcastScalarSupported1) {
-  TensorFactory<ScalarType::Float> tf;
-
-  Tensor a = tf.make({2, 1, 3}, {2, 3, 4, 5, 6, 7});
-  Tensor b = tf.make({1}, {2});
-
-  // Destination for the broadcasting div. Follow the broadcasting rules in
-  // https://fburl.com/n9wl4d0o
-  Tensor out = tf.zeros({2, 1, 3});
-
-  op_sub_out(a, b, 1, out);
-
-  Tensor ret = tf.make({2, 1, 3}, {0, 1, 2, 3, 4, 5});
-  EXPECT_TENSOR_EQ(out, ret);
+  test_broadcast_rank1_scalar<ScalarType::Float>();
+  test_broadcast_rank1_scalar<ScalarType::Half>();
 }
 
 TEST_F(OpSubOutTest, BroadcastScalarSupported2) {